Update app.py

app.py

@@ -8,7 +8,7 @@ from transformers import AutoModel, AutoTokenizer, Qwen2VLForConditionalGenerati
 from qwen_vl_utils import process_vision_info
 
 # --- CONFIGURATION ---
-print(f"🚀 Démarrage RAG Finance (Version Originale + Fix Hallucination)...")
+print(f"🚀 Démarrage RAG Finance (Version Originale + FILTRAGE STRICT)...")
 
 # --- 1. DONNÉES ---
 try:
@@ -18,9 +18,9 @@ except:
     dataset = []
     print("⚠️ Index vide ou fichier non trouvé.")
 
-# --- 2. MODÈLES (INCHANGÉS) ---
+# --- 2. MODÈLES (INCHANGÉS - ON GARDE LES GROS) ---
 
-# A. EMBEDDING : GTE-Qwen2-7B (Le modèle LOURD original)
+# A. EMBEDDING
 EMBED_MODEL_ID = "Alibaba-NLP/gte-Qwen2-7B-instruct"
 print(f"🔹 Chargement Embedder : {EMBED_MODEL_ID}")
 
@@ -29,8 +29,6 @@ embed_model = AutoModel.from_pretrained(
     EMBED_MODEL_ID,
     trust_remote_code=False,
     torch_dtype=torch.bfloat16,
-    # J'ai mis en commentaire la ligne qui fait planter le démarrage sur CPU (ZeroGPU)
-    # attn_implementation="flash_attention_2", 
     device_map="auto"
 )
 
@@ -50,13 +48,11 @@ print(f"👁️ Chargement Vision : {GEN_MODEL_ID}")
 gen_model = Qwen2VLForConditionalGeneration.from_pretrained(
     GEN_MODEL_ID, 
     torch_dtype=torch.bfloat16, 
-    # Idem, désactivé pour éviter le crash "No CUDA" au boot
-    # attn_implementation="flash_attention_2", 
     device_map="auto"
 )
 gen_processor = AutoProcessor.from_pretrained(GEN_MODEL_ID)
 
-# --- 3. FONCTIONS UTILITAIRES ---
+# --- 3. FONCTIONS ---
 
 def last_token_pool(last_hidden_states, attention_mask):
     left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
@@ -67,7 +63,7 @@ def last_token_pool(last_hidden_states, attention_mask):
         batch_size = last_hidden_states.shape[0]
         return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]
 
-# --- 4. LOGIQUE RAG MULTI-VIEW ---
+# --- 4. PIPELINE ---
 
 @spaces.GPU
 def retrieve_and_answer(query):
@@ -75,13 +71,38 @@ def retrieve_and_answer(query):
     
     if not dataset: return None, "Base vide", "Pas de document"
     
-    # === ÉTAPE 1 : RETRIEVAL (Embedding) ===
+    # === ÉTAPE 0 : FILTRAGE STRICT (LA SÉCURITÉ) ===
+    # On regarde si l'utilisateur parle d'une entreprise spécifique
+    # Si oui, on retire TOUTES les autres pages de la liste.
+    
+    query_lower = query.lower()
+    target_company = None
+    
+    if "microsoft" in query_lower or "msft" in query_lower:
+        target_company = "Microsoft"
+    elif "apple" in query_lower or "aapl" in query_lower:
+        target_company = "Apple"
+    elif "tesla" in query_lower:
+        # Cas piège Tesla : On sait qu'on n'a pas les docs, on coupe tout de suite.
+        return [], "", "Data not found: No documents available for Tesla in the database."
+
+    # === ÉTAPE 1 : RETRIEVAL ===
     valid_docs = []
     for i, doc in enumerate(dataset):
+        # Le Filtrage Strict s'applique ici
+        if target_company:
+            # Si on cherche Microsoft, on ignore tout ce qui ne contient pas "Microsoft" dans le nom du doc
+            if target_company not in doc.get('doc_name', ''):
+                continue 
+        
         text = doc.get('text', '').strip()
         if text:
-            valid_docs.append({'text': text, 'original_index': i})
+            valid_docs.append({'text': text, 'original_index': i, 'doc_name': doc.get('doc_name', 'Doc')})
             
+    # Si après filtrage on a plus rien (ex: question sur Tesla mal gérée avant), on arrête
+    if not valid_docs:
+        return [], "", "No relevant documents found for this company."
+
     query_text = f"Instruct: Given a user query, retrieve relevant passages that answer the query.\nQuery: {query}"
     
     with torch.no_grad():
@@ -93,7 +114,7 @@ def retrieve_and_answer(query):
         d_embeddings_list = []
         doc_texts = [d['text'] for d in valid_docs]
         
-        # Batch size de 1 pour économiser la mémoire avec le gros modèle 7B
+        # Batch size 1 pour le gros modèle 7B
         for i in range(0, len(doc_texts), 1):
             d_inputs = embed_tokenizer(doc_texts[i:i+1], max_length=8192, padding=True, truncation=True, return_tensors='pt').to(embed_model.device)
             d_outputs = embed_model(**d_inputs)
@@ -103,7 +124,9 @@ def retrieve_and_answer(query):
             
         d_emb_final = torch.cat(d_embeddings_list, dim=0)
         scores = (q_emb @ d_emb_final.T).squeeze(0)
-        top_k_indices = torch.topk(scores, k=min(10, len(scores))).indices.tolist()
+        # On prend max 10 ou moins si on a filtré
+        k_val = min(10, len(scores))
+        top_k_indices = torch.topk(scores, k=k_val).indices.tolist()
 
     # === ÉTAPE 2 : RERANKING ===
     pairs = []
@@ -113,9 +136,11 @@ def retrieve_and_answer(query):
     with torch.no_grad():
         r_inputs = rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=8192).to(rerank_model.device)
         r_scores = rerank_model(**r_inputs, return_dict=True).logits.view(-1).float()
-        top_3_indices_local = torch.topk(r_scores, k=min(3, len(r_scores))).indices.tolist()
+        
+        k_rerank = min(3, len(r_scores))
+        top_3_indices_local = torch.topk(r_scores, k=k_rerank).indices.tolist()
 
-    # === ÉTAPE 3 : PRÉPARATION IMAGES (ICI ON CORRIGE L'HALLUCINATION) ===
+    # === ÉTAPE 3 : PRÉPARATION IMAGES ===
     images_content = []
     gallery_preview = []
     meta_info = ""
@@ -128,31 +153,32 @@ def retrieve_and_answer(query):
         doc = dataset[final_doc_idx]
         image_path = doc['image_path']
         score = r_scores[idx_local].item()
+        doc_name = doc.get('doc_name', 'Unknown')
         
         try:
             img = Image.open(image_path)
             
-            # --- FIX HALLUCINATION ---
-            # On récupère le nom du document (ex: "Microsoft 2023 Report")
-            doc_name = doc.get('doc_name', 'Unknown Document')
-            
-            # On l'injecte explicitement dans le texte que voit l'IA
-            prompt_header = f"DOCUMENT SOURCE: {doc_name} (Relevance: {score:.2f})\n"
+            # Injection du nom pour aider encore plus
+            header_text = f"SOURCE DOCUMENT: {doc_name} (Confidence: {score:.2f})\n"
             
-            images_content.append({"type": "text", "text": prompt_header})
+            images_content.append({"type": "text", "text": header_text})
             images_content.append({"type": "image", "image": img})
             
-            gallery_preview.append((img, f"{doc_name} (Rank {rank+1})"))
+            gallery_preview.append((img, f"{doc_name}"))
             meta_info += f"- **{doc_name}** (Score: {score:.2f})\n"
         except:
             continue
 
-    # === ÉTAPE 4 : GÉNÉRATION ===
-    # On renforce le prompt système pour qu'il fasse attention au nom du document
+    if not images_content:
+        return [], "", "No images found."
+
+    # === 4. GENERATION ===
     system_prompt = (
-        "You are an expert financial analyst. Answer the user question using ONLY the provided images.\n"
-        "IMPORTANT: Before reading a table, check the 'DOCUMENT SOURCE' name above the image.\n"
-        "If the user asks about Microsoft, do not use data from an Apple document (and vice versa)."
+        "You are a strict financial analyst. Answer the user question using ONLY the provided images.\n"
+        "RULES:\n"
+        "1. If the user asks for 'Microsoft', ONLY use the image labeled 'Microsoft'. IGNORE Apple.\n"
+        "2. If the user asks for 'Apple', ONLY use the image labeled 'Apple'. IGNORE Microsoft.\n"
+        "3. Copy the exact number from the image. Do not calculate."
     )
     
     user_content = images_content + [{"type": "text", "text": f"\nUser Question: {query}"}]
@@ -172,7 +198,7 @@ def retrieve_and_answer(query):
         return_tensors="pt",
     ).to(gen_model.device)
     
-    generated_ids = gen_model.generate(**inputs, max_new_tokens=768)
+    generated_ids = gen_model.generate(**inputs, max_new_tokens=512)
     generated_ids_trimmed = [out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
     response = gen_processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
 
@@ -180,10 +206,10 @@ def retrieve_and_answer(query):
 
 # --- 5. UI ---
 with gr.Blocks(title="RAG Finance") as demo:
-    gr.Markdown("# 🚀 RAG Finance (Moteurs Originaux + Sécurité Hallucination)")
+    gr.Markdown("# 🚀 RAG Finance (Version Sécurisée)")
     
     with gr.Row():
-        query_input = gr.Textbox(label="Question", placeholder="Ex: What is the revenue of Microsoft?")
+        query_input = gr.Textbox(label="Question", placeholder="Ex: What is the Operating Income for Microsoft?")
         submit_btn = gr.Button("Analyser", variant="primary")
         
     with gr.Row():