Update app.py

app.py

@@ -1,3 +1,8 @@
+import os
+
+# جلوگیری از مشکل spaces watchdog
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
+
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -6,32 +11,55 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 MODEL_ID = "ddfws/Rezaeian-StatsAI"
 
 
-print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer = None
+model = None
 
 
-print("Loading model...")
+def load_model():
 
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto",
-    dtype=torch.float16,
-    trust_remote_code=True
-)
+    global tokenizer, model
+
+    if model is None:
 
-model.eval()
+        print("Loading tokenizer...")
+
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID
+        )
 
-print("Model loaded!")
 
+        print("Loading model...")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            device_map="auto",
+            dtype=torch.float16,
+            trust_remote_code=True
+        )
 
-def chat(message, history):
+        model.eval()
+
+        print("Model loaded!")
+
+
+
+def generate(message, history):
+
+    load_model()
 
     prompt = ""
 
-    for user, bot in history:
-        prompt += f"User: {user}\nAssistant: {bot}\n"
+    for h in history:
+        prompt += (
+            "User: " + h[0] +
+            "\nAssistant: " + h[1] +
+            "\n"
+        )
 
-    prompt += f"User: {message}\nAssistant:"
+    prompt += (
+        "User: " + message +
+        "\nAssistant:"
+    )
 
 
     inputs = tokenizer(
@@ -44,30 +72,31 @@ def chat(message, history):
 
         output = model.generate(
             **inputs,
-            max_new_tokens=512,
+            max_new_tokens=256,
             temperature=0.7,
             top_p=0.9,
             do_sample=True
         )
 
 
-    result = tokenizer.decode(
+    text = tokenizer.decode(
         output[0],
         skip_special_tokens=True
     )
 
 
-    answer = result.split("Assistant:")[-1]
+    answer = text.split("Assistant:")[-1]
 
     return answer
 
 
 
 demo = gr.ChatInterface(
-    fn=chat,
+    fn=generate,
     title="Rezaeian StatsAI",
     description="AI Assistant"
 )
 
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()