Update app.py

app.py

@@ -1,53 +1,33 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from sympy import solve, symbols
-from fastapi import FastAPI
-import uvicorn
 
-
-
-
-
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
-model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
-
-
-
-# مدل‌های مختلف
-MODEL_GENERAL = "meta-llama/Llama-2-7b-chat-hf"
-MODEL_IRAN = "HooshvareLab/bert-fa-base-uncased"
-MODEL_MATH = None  # SymPy برای ریاضی
-
-# بارگذاری مدل‌ها
-tokenizer_general = AutoTokenizer.from_pretrained(MODEL_GENERAL)
-model_general = AutoModelForCausalLM.from_pretrained(MODEL_GENERAL)
-
-tokenizer_iran = AutoTokenizer.from_pretrained(MODEL_IRAN)
-model_iran = AutoModelForCausalLM.from_pretrained(MODEL_IRAN)
-
-# FastAPI برای مدیریت درخواست‌ها
-app = FastAPI()
-
-def generate_response(model, tokenizer, prompt, max_tokens=100):
+# مدل‌های مختلف AI
+models = {
+    'falcon': "huggingface/falcon-7b-instruct",
+    'gpt_neox': "EleutherAI/gpt-neox-20b",
+    'persian': "HooshvareLab/bert-fa-zwnj-base",
+    'math': "EleutherAI/gpt-neox-20b-math"  # مدل ریاضی (باید ایجاد شود یا از مدل‌های مشابه استفاده کنید)
+}
+
+# بارگذاری مدل‌ها از Hugging Face
+tokenizers = {name: AutoTokenizer.from_pretrained(path) for name, path in models.items()}
+models = {name: AutoModelForCausalLM.from_pretrained(path) for name, path in models.items()}
+
+def generate_response(prompt, model_name="falcon"):
+    tokenizer = tokenizers[model_name]
+    model = models[model_name]
+    
     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)
-
-@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__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    outputs = model.generate(inputs["input_ids"], max_length=100, num_return_sequences=1)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return response
+
+# نمونه استفاده
+prompt = "سلام، امروز چه خبر؟"
+response_falcon = generate_response(prompt, model_name="falcon")
+response_gpt_neox = generate_response(prompt, model_name="gpt_neox")
+response_persian = generate_response(prompt, model_name="persian")
+
+print("Falcon Response:", response_falcon)
+print("GPT-NeoX Response:", response_gpt_neox)
+print("Persian Response:", response_persian)