Update app.py

app.py

@@ -1,85 +1,43 @@
-import spaces
-
 import gradio as gr
 import torch
-
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM
-)
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 MODEL_ID = "ddfws/Rezaeian-StatsAI"
 
 
 print("Loading tokenizer...")
-
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID
-)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 
 print("Loading model...")
 
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
-    torch_dtype=torch.float16
+    dtype=torch.float16,
+    trust_remote_code=True
 )
 
+model.eval()
 
 print("Model loaded!")
 
 
 def chat(message, history):
 
-    messages = []
-
-    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
-        }
-    )
+    prompt = ""
 
+    for user, bot in history:
+        prompt += f"User: {user}\nAssistant: {bot}\n"
 
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    prompt += f"User: {message}\nAssistant:"
 
 
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
-    )
-
-
-    if torch.cuda.is_available():
-        inputs = inputs.to("cuda")
+    ).to(model.device)
 
 
     with torch.no_grad():
@@ -94,28 +52,21 @@ def chat(message, history):
 
 
     result = tokenizer.decode(
-        output[0][inputs.input_ids.shape[1]:],
+        output[0],
         skip_special_tokens=True
     )
 
 
-    return result
+    answer = result.split("Assistant:")[-1]
+
+    return answer
 
 
 
 demo = gr.ChatInterface(
     fn=chat,
-    title="Rezaeian-StatsAI",
-    description="""
-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
-"""
+    title="Rezaeian StatsAI",
+    description="AI Assistant"
 )