Update app.py

app.py

@@ -63,19 +63,29 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
 
+model_id = "Asit03/AI_Agent_V2_Merged"
 
-model_name = "Asit03/AI_Agent_V2_Merged"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-def generate_response(prompt):
+# Load model with 4-bit quantization
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    load_in_4bit=True,
+    torch_dtype=torch.bfloat16,  # fallback to torch.float16 if needed
+    trust_remote_code=True
+)
+
+# Generation function
+def chat(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=1000)
+    outputs = model.generate(**inputs, max_new_tokens=200, pad_token_id=tokenizer.eos_token_id)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-gr.Interface(fn=generate_response, inputs="text", outputs="text").launch()
-
+# Launch Gradio app
+gr.Interface(fn=chat, inputs="text", outputs="text", title="💬 AI Agent V2").launch()