Update app.py

app.py

@@ -3,60 +3,66 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import gradio as gr
 
-# Your LoRA
+# Direct base model (no auto-detection needed)
+BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
 PEFT_ID = "rishu834763/java-explainer-lora"
 
-# Load base model in 8-bit instead of 4-bit (much faster & more stable cold start on free tier)
-base_model = "mistralai/Mistral-7B-Instruct-v0.2"
-
+# Load in 8-bit → super stable & fast cold start on free tier
 model = AutoModelForCausalLM.from_pretrained(
-    base_model,
+    BASE_MODEL,
     device_map="auto",
-    load_in_8bit=True,          # ← 8-bit instead of 4-bit = instant start
+    load_in_8bit=True,
     torch_dtype=torch.bfloat16,
 )
 
-# Apply your LoRA (no merge = super fast)
+# Load your LoRA on top (no merge = instant)
 model = PeftModel.from_pretrained(model, PEFT_ID)
 
-tokenizer = AutoTokenizer.from_pretrained(base_model)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+tokenizer.pad_token = tokenizer.eos_token
 
-# Proper generation function
 def chat(message, history):
     messages = []
-    for user, assistant in history:
-        messages.append({"role": "user", "content": user})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
-    
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
+
     output_ids = model.generate(
         input_ids,
         max_new_tokens=1024,
+        do_sample=True,
         temperature=0.6,
         top_p=0.9,
-        do_sample=True,
-        repetition_penalty=1.1,
-        eos_token_id=tokenizer.eos_token_id,
+        eos_token_id=terminators,
         pad_token_id=tokenizer.eos_token_id,
     )
-    
-    response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+
+    response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
     return response
 
-# Interface
+# Minimal interface – starts instantly
 gr.ChatInterface(
     chat,
-    title="Java Explainer ⚡ Live Now",
-    description="Your own fine-tuned model · Starts in seconds · No OpenAI",
+    title="Java Explainer – Your Model",
     examples=[
-        "Explain this Java code simply: public static void main(String[] args) { System.out.println(\"Hello\"); }",
-        "What is the difference between == and .equals()?",
-        "Why do we need the 'static' keyword in main()?",
+        "Explain this Java code simply:\npublic class Test {\n    public static void main(String[] args) {\n        System.out.println(\"Hello\");\n    }\n}",
+        "What is the difference between HashMap and Hashtable?",
+        "Why main method is public static void?",
     ],
     cache_examples=False,
-).queue(max_size=30).launch()
+    submit_btn="Send",
+    retry_btn=None,   # removes the retry that sometimes causes blank replies
+).queue(max_size=20).launch()