Update app.py

app.py

@@ -3,6 +3,10 @@ import os
 from transformers import pipeline
 from huggingface_hub import HfFolder
 
+import transformers
+import torch
+
+
 # Ensure the HF_TOKEN environment variable is set correctly
 HF_TOKEN = os.getenv('HF_TOKEN')
 if HF_TOKEN:
@@ -11,7 +15,15 @@ else:
     st.warning("HF_TOKEN is not set. Proceeding without a token.")
 
 # Use a valid model identifier
-generator = pipeline("text-generation", model="gpt2")
+#generator = pipeline("text-generation", model="gpt2")
+
+
+model_id = "meta-llama/Meta-Llama-3-8B"
+
+generator = transformers.pipeline(
+    "text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto"
+)
+
 
 st.title("Text Generation")
 st.write("Enter your text below.")