Update app.py

app.py

@@ -3,6 +3,8 @@ import requests
 import pandas as pd
 import os
 import time
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from templates.prompt_mistral_rag import RAG_PROMPT_TEMPLATE
 
 
 # URL du backend Hugging Face Space
@@ -10,6 +12,14 @@ API_URL = os.getenv('API_URL')
 API_URL_ASK =  API_URL+"/get_answer"
 SPACE_URL = "https://huggingface.co/api/spaces/Loren/api_search_articles"
 
+# Chargement du modèle génératif
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME,
+                                             torch_dtype=torch.float16,
+                                             device_map="auto"
+                                            )
+
 def get_tags():
     resp = requests.get(f"{API_URL}/get_tags")
     if resp.status_code != 200:
@@ -178,16 +188,9 @@ def get_answer_with_query(query, use_rerank, history):
     if not query:
         raise gr.Error("❌ Erreur : aucun query fourni.")
         return None
-    payload = {
-        "question": query,
-        "use_rerank": use_rerank
-    }
-    headers = {
-        "Content-Type": "application/json"
-    }
-    resp = requests.post(API_URL_ASK, json=payload, headers=headers)
-
-    # Vérification du statut HTTP
+    params = {"query": query, "use_rerank": use_rerank}
+
+    resp = requests.get(f"{API_URL}/get_query_results", params=params)
     if resp.status_code != 200:
         raise gr.Error(f"❌ Erreur : {resp.status_code}")
         return None
@@ -196,6 +199,23 @@ def get_answer_with_query(query, use_rerank, history):
         raise gr.Error(f"❌ Erreur : {dict_resp['code']} - {dict_resp['message']}")
         return None
 
+    list_chunks = [resp['chunk_text'] for resp in dict_resp['result']]
+    if not list_chunks:
+        answer = ("Je ne dispose pas d’informations sur ce sujet. "
+                  "Je peux uniquement répondre à des questions sur les articles " \
+                  "du jeu de données.")
+    else:
+        # Construction du prompt
+        prompt = RAG_PROMPT_TEMPLATE.format(
+            context="\n".join(list_chunks),
+            question=user_query
+        )
+        # Génération de la réponse
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        outputs = model.generate(**inputs, max_new_tokens=500)
+        generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]  # uniquement la partie générée
+        answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+
     df = pd.DataFrame(dict_resp["results"])
     # Convertir les URLs en liens HTML cliquables
     df["article_url"] = df["article_url"].apply(lambda x: f'<a href="{x}" target="_blank">Ouvrir</a>')
@@ -244,6 +264,7 @@ def get_answer_with_query(query, use_rerank, history):
         {html}
     </div>
     """
+    
     history.append((query, dict_resp['answer']))
     return "", history, styled_html