Update app.py

app.py

@@ -20,7 +20,8 @@ TEXT_MODEL_REPO = os.environ.get("TEXT_MODEL_REPO", "<your-username>/text-ft-foo
 CLIP_MODEL_REPO = os.environ.get("CLIP_MODEL_REPO", "<your-username>/clip-ft-food-rag")
 DATASET_REPO    = os.environ.get("DATASET_REPO",    "<your-username>/food-rag-index")
 # LLM via Inference API (set HF_TOKEN in Space secrets). Change to your preferred instruct model.
-LLM_ID = os.environ.get("LLM_ID", "mistralai/Mistral-7B-Instruct-v0.3")
+LLM_ID = os.environ.get("LLM_ID", "google/gemma-2-2b-it")
+
 # =================================================
 
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -105,11 +106,22 @@ def build_prompt(question: str, ctx: List[Pair]) -> str:
     return "\n".join(lines)
 
 def call_llm(prompt: str) -> str:
+    # prompt already includes your Context + Question text
     if client is None:
         return "(LLM not configured)\n\n" + prompt
     try:
-        out = client.text_generation(prompt=prompt, max_new_tokens=256, temperature=0.2, do_sample=True)
-        return out.strip()
+        resp = client.chat_completion(
+            messages=[
+                {"role": "system", "content": (
+                    "You are a helpful assistant. Use the provided context to answer in Persian language; "
+                    "if it's not in the context, say you don't know."
+                )},
+                {"role": "user", "content": prompt},
+            ],
+            max_tokens=256,
+            temperature=0.2,
+        )
+        return resp.choices[0].message.content.strip()
     except Exception as e:
         return f"(LLM error: {e})\n\n" + prompt