Update app.py

app.py

@@ -4,19 +4,18 @@ from peft import PeftModel
 import torch
 import os
 
-# Hugging Face token (optional)
+# Optional Hugging Face token
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-BASE = "microsoft/phi-2"
-LORA = "abhi9953/abhi-ai"
+BASE_MODEL = "microsoft/phi-2"
+LORA_MODEL = "abhi9953/abhi-ai"
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(BASE)
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 tokenizer.pad_token = tokenizer.eos_token
 
-# Load base model (CPU optimized)
 model = AutoModelForCausalLM.from_pretrained(
-    BASE,
+    BASE_MODEL,
     torch_dtype=torch.float32,
     device_map="cpu"
 )
@@ -24,38 +23,45 @@ model = AutoModelForCausalLM.from_pretrained(
 # Load LoRA adapter
 model = PeftModel.from_pretrained(
     model,
-    LORA,
+    LORA_MODEL,
     token=HF_TOKEN
 )
 
 model.eval()
 
-# 🔥 WARM-UP (removes first-response lag)
 with torch.no_grad():
     _ = model.generate(
         **tokenizer("Hello", return_tensors="pt"),
-        max_new_tokens=5,
+        max_new_tokens=10,
         use_cache=True
     )
 
 def chat(msg):
-    prompt = f"### User:\n{msg}\n\n### Abhi:"
+    prompt = f"""### User:
+{msg}
+
+### Abhi:
+"""
+
     inputs = tokenizer(prompt, return_tensors="pt")
 
     with torch.no_grad():
         output = model.generate(
             **inputs,
-            max_new_tokens=32,   # 🔥 VERY IMPORTANT
-            do_sample=False,     # 🔥 MUCH faster
-            use_cache=True
+            max_new_tokens=128,    
+            do_sample=False,       
+            use_cache=True,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id
         )
 
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return text.split("### Abhi:")[-1].strip()
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    return decoded.split("### Abhi:")[-1].strip()
 
 gr.Interface(
     fn=chat,
-    inputs="text",
+    inputs=gr.Textbox(lines=3, placeholder="Talk to Abhi AI..."),
     outputs="text",
-    title="Abhi AI (Fast Mode)"
+    title="Abhi AI (Fast + Stable Mode)"
 ).launch()