Update app.py

app.py

@@ -17,11 +17,12 @@ try:
         trust_remote_code=True  # <-- ADDED HERE
     )
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        device_map="auto",
-        torch_dtype="auto",
-        trust_remote_code=True,  # <-- ADDED HERE
-        # load_in_4bit=True  # Uncomment for quantization
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype="auto",
+    trust_remote_code=True,
+    use_flash_attention_2=True,  # <-- Add this line
+    # load_in_4bit=True
     )
 except Exception as e:
     raise gr.Error(f"Error loading model: {str(e)}")