Update app.py

app.py

@@ -1,62 +1,44 @@
-import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from sympy import solve, symbols
+from fastapi import FastAPI
+import uvicorn
 
-# بارگذاری مدل و توکنایزر از Hugging Face
-MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"  # مدل Llama 2 (نسخه Chat)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
+# مدل‌های مختلف
+MODEL_GENERAL = "meta-llama/Llama-2-7b-chat-hf"
+MODEL_IRAN = "HooshvareLab/bert-fa-base-uncased"
+MODEL_MATH = None  # SymPy برای ریاضی
 
-# تعریف تابع پاسخ‌دهی
-def respond(
-    message: str,
-    history: list[tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    # ساختن prompt از پیام‌های قبلی
-    context = f"{system_message}\n"
-    for user_message, bot_response in history:
-        context += f"User: {user_message}\nBot: {bot_response}\n"
-    context += f"User: {message}\nBot:"
+# بارگذاری مدل‌ها
+tokenizer_general = AutoTokenizer.from_pretrained(MODEL_GENERAL)
+model_general = AutoModelForCausalLM.from_pretrained(MODEL_GENERAL)
 
-    # تولید پاسخ
-    inputs = tokenizer(context, return_tensors="pt", padding=True, truncation=True)
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response = response.split("Bot:")[-1].strip()  # استخراج پاسخ
+tokenizer_iran = AutoTokenizer.from_pretrained(MODEL_IRAN)
+model_iran = AutoModelForCausalLM.from_pretrained(MODEL_IRAN)
 
-    yield response
+# FastAPI برای مدیریت درخواست‌ها
+app = FastAPI()
 
+def generate_response(model, tokenizer, prompt, max_tokens=100):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(inputs.input_ids, max_new_tokens=max_tokens)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# رابط کاربری Gradio
-demo = gr.ChatInterface(
-    fn=respond,
-    additional_inputs=[
-        gr.Textbox(
-            value="You are an advanced and friendly assistant.",
-            label="System message",
-        ),
-        gr.Slider(
-            minimum=10, maximum=1024, value=256, step=1, label="Max new tokens"
-        ),
-        gr.Slider(
-            minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"
-        ),
-        gr.Slider(
-            minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"
-        ),
-    ],
-    title="Advanced Chatbot with Llama 2",
-    description="A conversational AI based on Llama 2 fine-tuned for chat.",
-)
+@app.post("/chat")
+def chat(input_text: str, mode: str = "general"):
+    if mode == "general":
+        response = generate_response(model_general, tokenizer_general, input_text)
+    elif mode == "iran":
+        response = generate_response(model_iran, tokenizer_iran, input_text)
+    elif mode == "math":
+        x = symbols("x")
+        try:
+            solution = solve(input_text, x)
+            response = f"Solution: {solution}"
+        except Exception as e:
+            response = f"Math error: {str(e)}"
+    else:
+        response = "Invalid mode selected."
+    return {"response": response}
 
 if __name__ == "__main__":
-    demo.launch()
+    uvicorn.run(app, host="0.0.0.0", port=8000)