Fix base model

app.py

@@ -4,14 +4,13 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-BASE_MODEL = "meta-llama/Llama-3.2-3B"
+BASE_MODEL = "unsloth/llama-3.2-3b-bnb-4bit"
 ADAPTER_MODEL = "devNaam/vakilai-llama32-3b-v1"
 
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
 model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    torch_dtype=torch.float16,
     device_map="auto"
 )
 
@@ -23,22 +22,18 @@ def vakil_ai(prompt):
 
     output = model.generate(
         **inputs,
-        max_new_tokens=300,
-        temperature=0.7,
-        top_p=0.9
+        max_new_tokens=200,
+        temperature=0.7
     )
 
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return response
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
 demo = gr.Interface(
     fn=vakil_ai,
-    inputs=gr.Textbox(lines=4, placeholder="Ask your legal question..."),
+    inputs=gr.Textbox(lines=4),
     outputs="text",
-    title="AI Vakil – Legal Assistant",
-    description="VakilAI powered by Llama 3.2"
+    title="AI Vakil"
 )
 
 demo.launch()

requirements.txt

@@ -3,4 +3,5 @@ torch
 accelerate
 peft
 sentencepiece
-gradio
+gradio
+bitsandbytes