Update app.py

app.py

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