Update app.py

app.py

@@ -1,26 +1,24 @@
 import gradio as gr
-import spaces
 import torch
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
-MODEL_ID = "ddfws/Rezaeian-StatsAI"
+MODEL_PATH = "ddfws/Rezaeian-StatsAI"
 
 
 print("Loading tokenizer...")
 
 tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID
+    MODEL_PATH
 )
 
 
 print("Loading model...")
 
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
+    MODEL_PATH,
     device_map="auto",
-    dtype=torch.float16
+    dtype=torch.float16 if torch.cuda.is_available() else torch.float32
 )
 
 model.eval()
@@ -28,58 +26,220 @@ model.eval()
 print("MODEL READY")
 
 
-@spaces.GPU
+def format_text(text):
+
+    # فرمول های رایج آماری
+    replacements = {
+        "β0": r"$\beta_0$",
+        "β1": r"$\beta_1$",
+        "σ2": r"$\sigma^2$",
+        "R2": r"$R^2$",
+        "xbar": r"$\bar{x}$",
+        "mu": r"$\mu$",
+    }
+
+    for a,b in replacements.items():
+        text = text.replace(a,b)
+
+    return text
+
+
+
 def generate(message, history):
 
-    prompt = ""
+    messages=[]
 
-    for user, assistant in history:
-        prompt += f"User: {user}\nAssistant: {assistant}\n"
+    for user, bot in history:
 
+        messages.append(
+            {
+                "role":"user",
+                "content":user
+            }
+        )
 
-    prompt += f"User: {message}\nAssistant:"
+        messages.append(
+            {
+                "role":"assistant",
+                "content":bot
+            }
+        )
+
+
+    messages.append(
+        {
+            "role":"user",
+            "content":message
+        }
+    )
+
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
 
 
     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(
             **inputs,
-            max_new_tokens=256,
-            temperature=0.7,
+            max_new_tokens=700,
+            temperature=0.6,
             top_p=0.9,
-            do_sample=True
+            do_sample=True,
+            repetition_penalty=1.1
         )
 
 
-    result = tokenizer.decode(
-        output[0],
+
+    answer = tokenizer.decode(
+        output[0][inputs.input_ids.shape[1]:],
         skip_special_tokens=True
     )
 
 
-    return result.split("Assistant:")[-1]
+    answer = format_text(answer)
+
+
+    return answer
+
+
+
+
+css = """
+
+/* حذف موارد اضافی */
+
+footer {
+    display:none !important;
+}
+
+
+.gradio-container {
+
+    max-width:1100px !important;
+    direction:rtl;
+
+}
+
+
+/* پیام ها */
+
+.message {
+
+    direction:rtl !important;
+    text-align:right !important;
+    unicode-bidi:plaintext;
+
+}
+
 
+/* متن فارسی */
+
+.markdown {
+
+    direction:rtl;
+    text-align:right;
+
+}
+
+
+/* ورودی */
+
+textarea {
+
+    direction:rtl !important;
+    text-align:right !important;
+
+}
+
+
+/* فرمول ها */
+
+.katex {
+
+    direction:ltr !important;
+
+}
+
+
+.MathJax {
+
+    direction:ltr !important;
+
+}
+
+
+/* کد و انگلیسی */
+
+code, pre {
+
+    direction:ltr !important;
+    text-align:left !important;
+
+}
+
+"""
+
+
+head = """
+
+<script>
+window.MathJax = {
+ tex: {
+   inlineMath: [['$', '$']],
+   displayMath: [['$$','$$']]
+ }
+};
+</script>
+
+<script 
+src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js">
+</script>
+
+"""
 
 
 demo = gr.ChatInterface(
+
     fn=generate,
+
     title="Rezaeian StatsAI",
-    description="ZeroGPU AI Assistant"
+
+    description=
+    "دستیار هوش مصنوعی آمار و احتمال مهندسی",
+
+    css=css,
+
+    head=head,
+
+    chatbot=gr.Chatbot(
+        height=650,
+        show_copy_button=True
+    ),
+
+    examples=[
+        "فرمول میانگین نمونه را بنویس",
+        "رگرسیون خطی را با فرمول توضیح بده",
+        "آزمون t چیست؟",
+        "انحراف معیار را محاسبه کن"
+    ]
+
 )
 
 
+
 demo.launch(
-    server_name="0.0.0.0"
+    server_name="0.0.0.0",
+    server_port=7860,
+    show_api=False
 )