Update app.py

app.py

@@ -1,36 +1,31 @@
+import spaces
+
 import gradio as gr
 import torch
 
 from transformers import (
     AutoTokenizer,
-    AutoModelForCausalLM,
-    BitsAndBytesConfig
+    AutoModelForCausalLM
 )
 
 
-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 model...")
 
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_use_double_quant=True
-)
-
-
 model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    quantization_config=bnb_config,
-    device_map="auto"
+    MODEL_ID,
+    device_map="auto",
+    torch_dtype=torch.float16
 )
 
 
@@ -41,20 +36,36 @@ def chat(message, history):
 
     messages = []
 
-    for user, assistant in history:
-        messages.append(
-            {"role":"user","content":user}
-        )
-        messages.append(
-            {"role":"assistant","content":assistant}
-        )
+    for item in history:
+        if isinstance(item, dict):
+            messages.append(item)
+        else:
+            user_msg, bot_msg = item
+
+            messages.append(
+                {
+                    "role": "user",
+                    "content": user_msg
+                }
+            )
+
+            messages.append(
+                {
+                    "role": "assistant",
+                    "content": bot_msg
+                }
+            )
+
 
     messages.append(
-        {"role":"user","content":message}
+        {
+            "role": "user",
+            "content": message
+        }
     )
 
 
-    text = tokenizer.apply_chat_template(
+    prompt = tokenizer.apply_chat_template(
         messages,
         tokenize=False,
         add_generation_prompt=True
@@ -62,17 +73,24 @@ def chat(message, history):
 
 
     inputs = tokenizer(
-        text,
+        prompt,
         return_tensors="pt"
-    ).to(model.device)
+    )
 
 
-    output = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        temperature=0.7,
-        do_sample=True
-    )
+    if torch.cuda.is_available():
+        inputs = inputs.to("cuda")
+
+
+    with torch.no_grad():
+
+        output = model.generate(
+            **inputs,
+            max_new_tokens=512,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True
+        )
 
 
     result = tokenizer.decode(
@@ -89,9 +107,14 @@ demo = gr.ChatInterface(
     fn=chat,
     title="Rezaeian-StatsAI",
     description="""
-فارسی AI آماری ساخته شده توسط امیرحسین رضائیان.
-Fine-tuned on university statistics notes.
-Base model: Qwen2.5-7B-Instruct.
+Rezaeian-StatsAI
+
+Created by Amirhossein Rezaeian.
+
+A Persian statistical AI model fine-tuned on University of Tehran statistics notes.
+
+Base model:
+Qwen2.5-7B-Instruct
 """
 )