Update app.py

app.py

@@ -1,68 +1,105 @@
+# app.py
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 from peft import PeftModel
 import gradio as gr
 
-# Direct base model (no auto-detection needed)
-BASE_MODEL = "mistralai/Mistral-7B-Instruct-v0.2"
-PEFT_ID = "rishu834763/java-explainer-lora"
+# ===================================
+# 1. Model & LoRA (your exact repo)
+# ===================================
+BASE_MODEL = "meta-llama/Meta-Llama-3-8B-Instruct"   # do NOT change
+LORA_ADAPTER = "rishu834763/java-explainer-lora"     # ← your LoRA
 
-# Load in 8-bit → super stable & fast cold start on free tier
-model = AutoModelForCausalLM.from_pretrained(
+# 4-bit quantization (fits on 1×A100 40/80GB, 4090 24GB, T4 16GB with some offloading)
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+print("Loading base model (Llama-3-8B-Instruct 4-bit)...")
+base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    device_map="auto",
-    load_in_8bit=True,
+    quantization_config=quantization_config,
+    device_map="auto",           # auto-offload to CPU if needed
     torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
 )
 
-# Load your LoRA on top (no merge = instant)
-model = PeftModel.from_pretrained(model, PEFT_ID)
+print("Loading your LoRA adapter...")
+model = PeftModel.from_pretrained(base_model, LORA_ADAPTER)
 
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token
 
-def chat(message, history):
-    messages = []
-    for user_msg, bot_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        if bot_msg:
-            messages.append({"role": "assistant", "content": bot_msg})
+# ===================================
+# 2. Inference pipeline
+# ===================================
+pipe = torch.pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=1024,
+    temperature=0.3,
+    top_p=0.95,
+    do_sample=True,
+    repetition_penalty=1.15,
+    return_full_text=False,
+)
+
+# System prompt tuned for Java explanations
+SYSTEM_PROMPT = "You are an expert Java teacher. Explain concepts clearly, provide code examples, and answer in a concise but complete way."
+
+def chat(message: str, history):
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+
+    # Convert Gradio history → Llama-3 format
+    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})
 
-    input_ids = tokenizer.apply_chat_template(
+    prompt = tokenizer.apply_chat_template(
         messages,
+        tokenize=False,
         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,
-        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)
-    return response
-
-# Minimal interface – starts instantly
-gr.ChatInterface(
-    chat,
-    title="Java Explainer – Your Model",
-    examples=[
-        "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,
-    submit_btn="Send",
-    retry_btn=None,   # removes the retry that sometimes causes blank replies
-).queue(max_size=20).launch()
+    output = pipe(prompt)[0]["generated_text"]
+    return output
+
+# ===================================
+# 3. Modern Gradio UI (2025)
+# ===================================
+with gr.Blocks(theme=gr.themes.Soft(), title="Java Explainer (Llama-3-8B + Your LoRA)") as demo:
+    gr.Markdown("# 🧑‍💻 Java Explainer\nPowered by **rishu834763/java-explainer-lora** on Llama-3-8B-Instruct")
+
+    chatbot = gr.Chatbot(height=620)
+    msg = gr.Textbox(
+        placeholder="Ask anything about Java (e.g. 'Explain Spring Boot @Autowired with example')",
+        label="Your question",
+        container=False,
+    )
+
+    with gr.Row():
+        send = gr.Button("Send 🚀", variant="primary")
+        clear = gr.Button("Clear 🗑️")
+
+    with gr.Row():
+        retry = gr.Button("🔄 Retry")
+        undo = gr.Button("↶ Undo")
+
+    # Events
+    send.click(chat, [msg, chatbot], [msg, chatbot]).then(lambda: "", outputs=msg)
+    msg.submit(chat, [msg, chatbot], [msg, chatbot]).then(lambda: "", outputs=msg)
+    clear.click(lambda: None, None, chatbot, queue=False)
+    retry.click(lambda h: h[:-1], chatbot, chatbot, queue=False)
+    undo.click(lambda h: h[:-1], chatbot, chatbot, queue=False)
+
+demo.queue(max_size=64).launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+)