Update app.py

app.py

@@ -1,44 +1,59 @@
-# app.py
+import os
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-
-import os
-
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache"
-
-
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from llama_cpp import Llama
 
 app = FastAPI(
-    title="MGZON FLAN‑T5 API",
-    description="API for MGZON FLAN‑T5 model",
+    title="MGZON Smart Assistant",
+    description="دمج نموذج T5 المدرب مع Mistral 7B (GGUF) داخل Space"
 )
 
-MODEL_NAME = "MGZON/mgzon-flan-t5-base"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
-
-
-class PromptRequest(BaseModel):
-    prompt: str
-    max_new_tokens: int = 100
-
+# -------------------------------------------------
+# 1️⃣ تحميل نموذج T5 المدرب من Hub
+# -------------------------------------------------
+T5_REPO = "MGZON/mgzon-flan-t5-base"          # إذا رفعته باسم آخر عدّل هنا
+t5_tokenizer = AutoTokenizer.from_pretrained(T5_REPO)
+t5_model = AutoModelForSeq2SeqLM.from_pretrained(T5_REPO)
+
+# -------------------------------------------------
+# 2️⃣ تحميل ملف Mistral .gguf (تم تنزيله بـ setup.sh)
+# -------------------------------------------------
+gguf_path = os.path.abspath("models/mistral-7b-instruct-v0.1.Q4_K_M.gguf")
+mistral = Llama(
+    model_path=gguf_path,
+    n_ctx=2048,
+    n_threads=8,                 # عدد الـ threads المتاح داخل Space
+    # إذا كان لديك GPU داخل Space يمكنك إضافة n_gpu_layers=35
+)
 
-@app.post("/generate")
-def generate(req: PromptRequest):
-    """
-    استدعاء النموذج وإرجاع النص المولد.
-    """
-    try:
-        inputs = tokenizer(req.prompt, return_tensors="pt")
-        out_ids = model.generate(
-            **inputs,
-            max_new_tokens=req.max_new_tokens,
-            num_beams=4,
-            early_stopping=True,
-        )
-        text = tokenizer.decode(out_ids[0], skip_special_tokens=True)
-        return {"generated_text": text}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+# -------------------------------------------------
+# 3️⃣ تعريف شكل الطلب
+# -------------------------------------------------
+class AskRequest(BaseModel):
+    question: str
+    max_new_tokens: int = 150   # عدد الكلمات التي تريد استخراجها
+
+# -------------------------------------------------
+# 4️⃣ نقطة النهاية /ask
+# -------------------------------------------------
+@app.post("/ask")
+def ask(req: AskRequest):
+    q = req.question.strip()
+    if not q:
+        raise HTTPException(status_code=400, detail="Empty question")
+
+    # منطق بسيط لاختيار النموذج:
+    if any(tok in q.lower() for tok in ["mgzon", "flan", "t5"]):
+        # ----- استخدام T5 -----
+        inputs = t5_tokenizer(q, return_tensors="pt", truncation=True, max_length=256)
+        out_ids = t5_model.generate(**inputs, max_length=req.max_new_tokens)
+        answer = t5_tokenizer.decode(out_ids[0], skip_special_tokens=True)
+        model_name = "MGZON‑FLAN‑T5"
+    else:
+        # ----- استخدام Mistral -----
+        out = mistral(prompt=q, max_tokens=req.max_new_tokens)
+        answer = out["choices"][0]["text"].strip()
+        model_name = "Mistral‑7B‑GGUF"
+
+    return {"model": model_name, "response": answer}