Upgrade VakilAI chat interface

app.py

@@ -4,36 +4,76 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
+# Base model used during training
 BASE_MODEL = "unsloth/llama-3.2-3b-bnb-4bit"
+
+# Your VakilAI LoRA adapter
 ADAPTER_MODEL = "devNaam/vakilai-llama32-3b-v1"
 
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
+print("Loading base model...")
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
     device_map="auto"
 )
 
+print("Loading VakilAI adapter...")
 model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
 
-def vakil_ai(prompt):
+print("Model ready.")
+
+
+# Prompt template for legal assistant behavior
+def build_prompt(user_question):
+    return f"""
+You are VakilAI, an AI legal assistant specializing in Indian law.
+
+Explain legal concepts clearly in simple language.
+If possible, mention relevant IPC sections or legal principles.
+
+User Question:
+{user_question}
+
+Answer:
+"""
+
+
+# Generate response
+def vakil_ai(user_message, history):
+
+    prompt = build_prompt(user_message)
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     output = model.generate(
         **inputs,
-        max_new_tokens=200,
-        temperature=0.7
+        max_new_tokens=250,
+        temperature=0.5,
+        top_p=0.9
     )
 
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Extract only the answer part
+    if "Answer:" in response:
+        response = response.split("Answer:")[-1].strip()
+
+    return response
 
 
-demo = gr.Interface(
+# Chat interface
+demo = gr.ChatInterface(
     fn=vakil_ai,
-    inputs=gr.Textbox(lines=4),
-    outputs="text",
-    title="AI Vakil"
+    title="⚖️ AI Vakil – Indian Legal Assistant",
+    description="Ask questions about Indian law, IPC sections, and legal concepts.",
+    examples=[
+        "What is IPC Section 307?",
+        "What is the punishment for theft in India?",
+        "What is the difference between murder and culpable homicide?",
+        "What rights does a person have during arrest in India?"
+    ]
 )
 
 demo.launch()