Update app.py

app.py

@@ -3,7 +3,8 @@ import torch
 
 from transformers import (
     AutoTokenizer,
-    AutoModelForCausalLM
+    AutoModelForCausalLM,
+    BitsAndBytesConfig
 )
 
 
@@ -11,30 +12,45 @@ MODEL_ID = "ddfws/Rezaeian-StatsAI"
 
 
 print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID
+)
+
+
+print("Loading 4bit model...")
+
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True
+)
 
-print("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
+    quantization_config=bnb_config,
     device_map="auto"
 )
 
 
 model.eval()
 
+print("Model loaded successfully")
+
 
 system_prompt = """
 You are Rezaeian-StatsAI.
-You are a university statistics and probability assistant.
-Answer in Persian.
-Solve engineering statistics problems step by step.
+You are a Persian university assistant specialized in Engineering Statistics and Probability.
+Answer questions in Persian.
+Solve problems step by step.
+Use clear mathematical explanations.
 """
 
 
-def chat(question):
+def answer(question):
 
     prompt = f"""
 <|im_start|>system
@@ -57,33 +73,40 @@ def chat(question):
 
     with torch.no_grad():
 
-        output = model.generate(
+        outputs = model.generate(
             **inputs,
             max_new_tokens=512,
             temperature=0.7,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
         )
 
 
-    answer = tokenizer.decode(
-        output[0],
+    text = tokenizer.decode(
+        outputs[0],
         skip_special_tokens=True
     )
 
 
-    return answer.split("assistant")[-1]
+    if "assistant" in text:
+        text = text.split("assistant")[-1]
+
+
+    return text.strip()
+
 
 
 demo = gr.Interface(
-    fn=chat,
+    fn=answer,
     inputs=gr.Textbox(
-        label="سوال آماری"
+        lines=3,
+        placeholder="سوال آماری خود را بنویسید..."
     ),
     outputs=gr.Textbox(
-        label="پاسخ مدل"
+        lines=10
     ),
     title="Rezaeian-StatsAI",
-    description="AI assistant for Engineering Statistics"
+    description="Persian Engineering Statistics AI"
 )