Update app.py

app.py

@@ -1,60 +1,62 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 import gradio as gr
 
+# Your LoRA
 PEFT_ID = "rishu834763/java-explainer-lora"
 
-# Get base model name
-config = PeftConfig.from_pretrained(PEFT_ID)
-base = config.base_model_name_or_path
+# 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 base model in 4-bit
 model = AutoModelForCausalLM.from_pretrained(
-    base,
+    base_model,
     device_map="auto",
+    load_in_8bit=True,          # ← 8-bit instead of 4-bit = instant start
     torch_dtype=torch.bfloat16,
-    load_in_4bit=True,
 )
 
-# Load LoRA weights on top — BUT DO NOT MERGE (this is the trick!)
+# Apply your LoRA (no merge = super fast)
 model = PeftModel.from_pretrained(model, PEFT_ID)
 
-# Tokenizer
-tokenizer = AutoTokenizer.from_pretrained(base)
+tokenizer = AutoTokenizer.from_pretrained(base_model)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-# Pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=1024,
-    temperature=0.6,
-    do_sample=True,
-    top_p=0.9,
-)
-
+# Proper generation function
 def chat(message, history):
     messages = []
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+    for user, assistant in history:
+        messages.append({"role": "user", "content": user})
+        if assistant:
+            messages.append({"role": "assistant", "content": assistant})
     messages.append({"role": "user", "content": message})
 
-    output = pipe(messages)[0]["generated_text"]
-    return output[-1]["content"]
-
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
+    
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=1024,
+        temperature=0.6,
+        top_p=0.9,
+        do_sample=True,
+        repetition_penalty=1.1,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    
+    response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
+    return response
+
+# Interface
 gr.ChatInterface(
     chat,
-    title="Java Explainer – Your Model (Running!)",
-    description="100% your fine-tuned LoRA · No OpenAI · Instant start",
+    title="Java Explainer ⚡ Live Now",
+    description="Your own fine-tuned model · Starts in seconds · No OpenAI",
     examples=[
-        "Explain this Java code: public static void main(String[] args) { System.out.println(\"Hello\"); }",
-        "What does public static void main mean?",
-        "Difference between String and StringBuilder?",
+        "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()?",
     ],
     cache_examples=False,
-).queue().launch()
+).queue(max_size=30).launch()