diff --git a/rag.py b/rag.py index b0f3127..21a57ba 100644 --- a/rag.py +++ b/rag.py @@ -17,7 +17,7 @@ from FlagEmbedding import BGEM3FlagModel from rich import print ROOT = Path("Fiches") # dossier racine des fiches -K = 4 # nombre de passages remis au LLM +K = 10 # nombre de passages remis au LLM # ------------------ utilitaires de découpe identiques à l'index ------------- CHUNK, OVERLAP = 800, 100 # garder cohérent avec index.py @@ -80,8 +80,10 @@ try: D, I = idx.search(q_emb.astype("float32").reshape(1, -1), K) hits = I[0] + # réordonne pour mettre en tête les passages contenant “Seuil” + hits = sorted(hits, key=lambda i: "Seuil" not in docs[int(i)]) - context = "\n\n".join(docs[int(i)] for i in hits) + context = "\n\n".join(docs[int(i)] for i in hits[:K]) prompt = ( "Réponds en français, de façon précise, et uniquement à partir du contexte fourni. Si l'information n'est pas dans le contexte, réponds : 'Je ne sais pas'.\n" f"{context}\n{q}"