Upgrade VakilAI chat interface

app.py

@@ -4,76 +4,62 @@ 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...")
+print("Loading base model on CPU...")
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    device_map="auto"
+    device_map="cpu",
+    torch_dtype=torch.float32
 )
 
 print("Loading VakilAI adapter...")
 model = PeftModel.from_pretrained(model, ADAPTER_MODEL)
 
-print("Model ready.")
+print("Model ready")
 
 
-# Prompt template for legal assistant behavior
-def build_prompt(user_question):
+def build_prompt(question):
     return f"""
-You are VakilAI, an AI legal assistant specializing in Indian law.
+You are VakilAI, an AI legal assistant for Indian law.
 
-Explain legal concepts clearly in simple language.
-If possible, mention relevant IPC sections or legal principles.
+Explain the answer clearly and simply.
 
-User Question:
-{user_question}
+Question:
+{question}
 
 Answer:
 """
 
 
-# Generate response
-def vakil_ai(user_message, history):
+def vakil_ai(message, history):
 
-    prompt = build_prompt(user_message)
+    prompt = build_prompt(message)
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(prompt, return_tensors="pt")
 
     output = model.generate(
         **inputs,
-        max_new_tokens=250,
-        temperature=0.5,
-        top_p=0.9
+        max_new_tokens=200,
+        temperature=0.5
     )
 
     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
 
 
-# Chat interface
 demo = gr.ChatInterface(
     fn=vakil_ai,
     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?"
-    ]
+    description="Ask questions about IPC, Indian law, and legal concepts.",
 )
 
 demo.launch()