Update app.py

app.py

@@ -8,105 +8,91 @@ from transformers import (
 )
 
 
-MODEL_ID = "ddfws/Rezaeian-StatsAI"
+model_id = "ddfws/Rezaeian-StatsAI"
 
 
 print("Loading tokenizer...")
 
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID
-)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
-print("Loading 4bit model...")
+print("Loading model...")
 
 
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
     bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4",
     bnb_4bit_use_double_quant=True
 )
 
 
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
+    model_id,
     quantization_config=bnb_config,
     device_map="auto"
 )
 
 
-model.eval()
-
-print("Model loaded successfully")
+print("Model loaded!")
 
 
-system_prompt = """
-You are Rezaeian-StatsAI.
-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(message, history):
 
+    messages = []
 
-def answer(question):
+    for user, assistant in history:
+        messages.append(
+            {"role":"user","content":user}
+        )
+        messages.append(
+            {"role":"assistant","content":assistant}
+        )
 
-    prompt = f"""
-<|im_start|>system
-{system_prompt}
-<|im_end|>
+    messages.append(
+        {"role":"user","content":message}
+    )
 
-<|im_start|>user
-{question}
-<|im_end|>
 
-<|im_start|>assistant
-"""
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
 
 
     inputs = tokenizer(
-        prompt,
+        text,
         return_tensors="pt"
     ).to(model.device)
 
 
-    with torch.no_grad():
-
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    output = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        do_sample=True
+    )
 
 
-    text = tokenizer.decode(
-        outputs[0],
+    result = tokenizer.decode(
+        output[0][inputs.input_ids.shape[1]:],
         skip_special_tokens=True
     )
 
 
-    if "assistant" in text:
-        text = text.split("assistant")[-1]
-
+    return result
 
-    return text.strip()
 
 
-
-demo = gr.Interface(
-    fn=answer,
-    inputs=gr.Textbox(
-        lines=3,
-        placeholder="سوال آماری خود را بنویسید..."
-    ),
-    outputs=gr.Textbox(
-        lines=10
-    ),
+demo = gr.ChatInterface(
+    fn=chat,
     title="Rezaeian-StatsAI",
-    description="Persian Engineering Statistics AI"
+    description="""
+فارسی AI آماری ساخته شده توسط امیرحسین رضائیان.
+Fine-tuned on university statistics notes.
+Base model: Qwen2.5-7B-Instruct.
+"""
 )