Switched back to running on GPU

app.py

@@ -6,8 +6,7 @@ import torch
 @st.cache_resource
 def load_model_and_tokenizer(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    # Load the model in 8-bit quantization
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", load_in_8bit=True)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     return model, tokenizer
 
 model_8b, tokenizer_8b = load_model_and_tokenizer("huggyllama/llama-3.1-8b")
@@ -19,15 +18,13 @@ def generate_text(model, tokenizer, prompt, max_length=100):
         outputs = model.generate(**inputs, max_length=max_length, num_return_sequences=1)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-st.title("LLaMA-3.1-8B vs LLaMA-3.1-8B-Instruct Comparison (CPU Version)")
+st.title("LLaMA-3.1-8B vs LLaMA-3.1-8B-Instruct Comparison")
 
 prompt = st.text_area("Enter your prompt:", height=100)
 max_length = st.slider("Max output length:", min_value=50, max_value=500, value=100)
 
 if st.button("Generate"):
     if prompt:
-        st.warning("Generation may take several minutes. Please be patient.")
-        
         col1, col2 = st.columns(2)
         
         with col1: