Update app.py

app.py

@@ -1,22 +1,26 @@
 import gradio as gr
+import spaces
 import torch
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
-MODEL_PATH = "ddfws/Rezaeian-StatsAI"
+MODEL_ID = "ddfws/Rezaeian-StatsAI"
 
 
 print("Loading tokenizer...")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID
+)
 
 
 print("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_PATH,
+    MODEL_ID,
     device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    dtype=torch.float16
 )
 
 model.eval()
@@ -25,154 +29,130 @@ print("MODEL READY")
 
 
 
-def format_text(text):
-
-    replacements = {
-        "β0": r"$\beta_0$",
-        "β1": r"$\beta_1$",
-        "σ2": r"$\sigma^2$",
-        "R2": r"$R^2$",
-        "x̄": r"$\bar{x}$",
-        "mu": r"$\mu$"
-    }
+def format_answer(text):
 
-    for a,b in replacements.items():
-        text=text.replace(a,b)
+    # تبدیل چند فرمول رایج آماری به LaTeX
+    text = text.replace("β0", "$\\beta_0$")
+    text = text.replace("β1", "$\\beta_1$")
+    text = text.replace("R2", "$R^2$")
+    text = text.replace("σ2", "$\\sigma^2$")
+    text = text.replace("x̄", "$\\bar{x}$")
 
     return text
 
 
 
+@spaces.GPU
 def generate(message, history):
 
-    messages=[]
-
-    for user, bot in history:
-
-        messages.append({
-            "role":"user",
-            "content":user
-        })
+    prompt = ""
 
-        messages.append({
-            "role":"assistant",
-            "content":bot
-        })
+    for user, assistant in history:
 
-
-    messages.append({
-        "role":"user",
-        "content":message
-    })
+        prompt += (
+            f"User: {user}\n"
+            f"Assistant: {assistant}\n"
+        )
 
 
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    prompt += f"User: {message}\nAssistant:"
 
 
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
-    ).to(model.device)
+    )
+
+
+    inputs = {
+        k:v.to(model.device)
+        for k,v in inputs.items()
+    }
 
 
     with torch.no_grad():
 
-        output=model.generate(
+        output = model.generate(
+
             **inputs,
-            max_new_tokens=700,
-            temperature=0.6,
+
+            max_new_tokens=512,
+
+            temperature=0.7,
+
             top_p=0.9,
+
             do_sample=True
         )
 
 
-    answer=tokenizer.decode(
-        output[0][inputs.input_ids.shape[1]:],
+    result = tokenizer.decode(
+        output[0],
         skip_special_tokens=True
     )
 
 
-    return format_text(answer)
+    answer = result.split("Assistant:")[-1]
+
 
+    return format_answer(answer)
 
 
-css="""
 
-footer{
+css = """
+
+footer {
 display:none !important;
 }
 
 
-.gradio-container{
-direction:rtl;
+.gradio-container {
+
+direction:rtl !important;
+
 }
 
 
-.message{
+.message {
+
 direction:rtl !important;
+
 text-align:right !important;
-unicode-bidi:plaintext;
-}
 
+unicode-bidi:plaintext;
 
-.markdown{
-direction:rtl;
-text-align:right;
 }
 
 
-textarea{
+textarea {
+
 direction:rtl !important;
-text-align:right !important;
-}
 
+text-align:right !important;
 
-code,pre{
-direction:ltr !important;
-text-align:left !important;
 }
 
 
-"""
+code, pre {
 
+direction:ltr !important;
 
-head="""
+text-align:left !important;
 
-<script>
-window.MathJax={
-tex:{
-inlineMath:[['$','$']],
-displayMath:[['$$','$$']]
 }
-};
-</script>
-
-<script src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
 
 """
 
 
-chatbot = gr.Chatbot(
-    height=600
-)
-
-
 demo = gr.ChatInterface(
     fn=generate,
-    chatbot=chatbot,
     title="Rezaeian StatsAI",
-    description="دستیار هوش مصنوعی آمار و احنمال مهندسی",
-    css=css,
-    head=head
+    description="دستیار هوش مصنوعی آمار",
+    css=css
 )
 
 
+
 demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    show_api=False
+    server_name="0.0.0.0"
 )