Update chatpdf.py

chatpdf.py

@@ -1,8 +1,3 @@
-# -*- coding: utf-8 -*-
-"""
-@author:XuMing([email protected])
-@description: 
-"""
 import argparse
 import hashlib
 import os
@@ -18,6 +13,7 @@ from similarities import (
     EnsembleSimilarity,
     BertSimilarity,
     BM25Similarity,
+    TfidfSimilarity
 )
 from similarities.similarity import SimilarityABC
 from transformers import (
@@ -43,10 +39,9 @@ MODEL_CLASSES = {
     "auto": (AutoModelForCausalLM, AutoTokenizer),
 }
 
-PROMPT_TEMPLATE = """Basándose en la siguiente información conocida, responda a la pregunta del usuario de forma 
-concisa y profesional. Si no puede obtener una respuesta, diga «No se puede responder a la pregunta basándose en la 
-información conocida» o «No se proporciona suficiente información relevante», no está permitido añadir elementos 
-inventados en la respuesta, y ésta debe estar en Español.
+PROMPT_TEMPLATE = """Basándose únicamente en la información proporcionada a continuación, responda a las preguntas del usuario de manera concisa y profesional. 
+No se debe responder a preguntas relacionadas con sentimientos, emociones, temas personales o cualquier información que no esté explícitamente presente en el contenido proporcionado. 
+Si la pregunta se refiere a un artículo específico y no se encuentra en el contenido proporcionado, diga: "No se puede encontrar el artículo solicitado en la información conocida".
 
 Contenido conocido:
 {context_str}
@@ -55,8 +50,6 @@ Pregunta:
 {query_str}
 """
 
-
-
 class SentenceSplitter:
     def __init__(self, chunk_size: int = 250, chunk_overlap: int = 50):
         self.chunk_size = chunk_size
@@ -121,8 +114,7 @@ class SentenceSplitter:
         return overlapped_chunks
 
 
-
-class Rag:
+class ChatPDF:
     def __init__(
             self,
             similarity_model: SimilarityABC = None,
@@ -139,8 +131,8 @@ class Rag:
             rerank_model_name_or_path: str = None,
             enable_history: bool = False,
             num_expand_context_chunk: int = 2,
-            similarity_top_k: int = 10,
-            rerank_top_k: int = 3,
+            similarity_top_k: int = 15,
+            rerank_top_k: int = 5,
     ):
         """
         Init RAG model.
@@ -176,9 +168,11 @@ class Rag:
         if similarity_model is not None:
             self.sim_model = similarity_model
         else:
-            m1 = BertSimilarity(model_name_or_path="hiiamsid/sentence_similarity_spanish_es", device=self.device)
+            m1 = BertSimilarity(model_name_or_path="sentence-transformers/all-mpnet-base-v2", device=self.device)
             m2 = BM25Similarity()
-            default_sim_model = EnsembleSimilarity(similarities=[m1, m2], weights=[0.5, 0.5], c=2)
+            m3 = TfidfSimilarity()
+            default_sim_model = EnsembleSimilarity(similarities=[m1, m2, m3], weights=[0.5, 0.5, 0.5],
+                                                   c=2)  # Ajuste los pesos según los resultados
             self.sim_model = default_sim_model
         self.gen_model, self.tokenizer = self._init_gen_model(
             generate_model_type,
@@ -243,14 +237,14 @@ class Rag:
         try:
             model.generation_config = GenerationConfig.from_pretrained(gen_model_name_or_path, trust_remote_code=True)
         except Exception as e:
-            logger.warning(f"Failed to load generation config from {gen_model_name_or_path}, {e}")
+            logger.warning(f"No se pudo cargar la configuración de generación desde {gen_model_name_or_path}, {e}")
         if peft_name:
             model = PeftModel.from_pretrained(
                 model,
                 peft_name,
                 torch_dtype="auto",
             )
-            logger.info(f"Loaded peft model from {peft_name}")
+            logger.info(f"Modelo peft cargado desde {peft_name}")
         model.eval()
         return model, tokenizer
 
@@ -341,6 +335,7 @@ class Rag:
                 raw_text = [text.strip() for text in page_text.splitlines() if text.strip()]
                 new_text = ''
                 for text in raw_text:
+                    # Añadir un espacio antes de concatenar si new_text no está vacío
                     if new_text:
                         new_text += ' '
                     new_text += text
@@ -397,25 +392,37 @@ class Rag:
         return scores
 
     def get_reference_results(self, query: str):
-        """
-        Get reference results.
-            1. Similarity model get similar chunks
-            2. Rerank similar chunks
-            3. Expand reference context chunk
-        :param query:
-        :return:
-        """
-        reference_results = []
+        # Verificar si la consulta incluye un "Artículo X"
+        exact_match = None
+        if re.search(r'Artículo\s*\d+', query, re.IGNORECASE):
+            # Buscar el término específico "Artículo X" en el corpus de manera más precisa
+            term = re.search(r'Artículo\s*\d+', query, re.IGNORECASE).group()
+            # Buscar coincidencias exactas en el corpus
+            for corpus_id, content in self.sim_model.corpus.items():
+                # Agregar espacio o signo de puntuación alrededor de "term" para evitar coincidencias parciales
+                if re.search(r'\b' + re.escape(term) + r'\b', content, re.IGNORECASE):
+                    exact_match = content
+                    break
+
+        if exact_match:
+            # Si se encuentra una coincidencia exacta, devolverla como contexto
+            return [exact_match]
+
+        # Si no se encuentra una coincidencia exacta, continuar con la búsqueda general
         sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
-        # Get reference results from corpus
-        hit_chunk_dict = dict()
-        for c in sim_contents:
-            for id_score_dict in c:
-                corpus_id = id_score_dict['corpus_id']
-                hit_chunk = id_score_dict["corpus_doc"]
-                reference_results.append(hit_chunk)
-                hit_chunk_dict[corpus_id] = hit_chunk
 
+        # Procesar los resultados de similitud
+        reference_results = []
+
+        hit_chunk_dict = dict()
+        threshold_score = 0.5  # Establece un umbral para filtrar fragmentos irrelevantes
+
+        for query_id, id_score_dict in sim_contents.items():
+            for corpus_id, s in id_score_dict.items():
+                if s > threshold_score:  # Filtrar por puntuación de similitud
+                    hit_chunk = self.sim_model.corpus[corpus_id]
+                    reference_results.append(hit_chunk)
+                    hit_chunk_dict[corpus_id] = hit_chunk
         if reference_results:
             if self.rerank_model is not None:
                 # Rerank reference results
@@ -440,9 +447,9 @@ class Rag:
     def predict_stream(
             self,
             query: str,
-            max_length: int = 512,
-            context_len: int = 2048,
-            temperature: float = 0.7,
+            max_length: int = 256,
+            context_len: int = 1024,
+            temperature: float = 0.5,
     ):
         """Generate predictions stream."""
         stop_str = self.tokenizer.eos_token if self.tokenizer.eos_token else "</s>"
@@ -450,15 +457,16 @@ class Rag:
             self.history = []
         if self.sim_model.corpus:
             reference_results = self.get_reference_results(query)
-            if reference_results:
-                reference_results = self._add_source_numbers(reference_results)
-                context_str = '\n'.join(reference_results)[:(context_len - len(PROMPT_TEMPLATE))]
-            else:
-                context_str = ''
+            if not reference_results:
+                yield 'No se ha proporcionado suficiente información relevante', reference_results
+            reference_results = self._add_source_numbers(reference_results)
+            context_str = '\n'.join(reference_results)[:]
+            print("gggggg: ", (context_len - len(PROMPT_TEMPLATE)))
             prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
+            logger.debug(f"prompt: {prompt}")
         else:
             prompt = query
-        logger.debug(f"prompt: {prompt}")
+            logger.debug(prompt)
         self.history.append([prompt, ''])
         response = ""
         for new_text in self.stream_generate_answer(
@@ -473,9 +481,9 @@ class Rag:
     def predict(
             self,
             query: str,
-            max_length: int = 512,
-            context_len: int = 2048,
-            temperature: float = 0.7,
+            max_length: int = 256,
+            context_len: int = 1024,
+            temperature: float = 0.5,
     ):
         """Query from corpus."""
         reference_results = []
@@ -483,15 +491,20 @@ class Rag:
             self.history = []
         if self.sim_model.corpus:
             reference_results = self.get_reference_results(query)
-            if reference_results:
-                reference_results = self._add_source_numbers(reference_results)
-                context_str = '\n'.join(reference_results)[:(context_len - len(PROMPT_TEMPLATE))]
-            else:
-                context_str = ''
+
+            if not reference_results:
+                return 'No se ha proporcionado suficiente información relevante', reference_results
+            reference_results = self._add_source_numbers(reference_results)
+            # context_str = '\n'.join(reference_results)  # Usa todos los fragmentos
+            context_st = '\n'.join(reference_results)[:(context_len - len(PROMPT_TEMPLATE))]
+            #print("Context: ", (context_len - len(PROMPT_TEMPLATE)))
+            print(".......................................................")
+            context_str = '\n'.join(reference_results)[:]
+            #print("context_str: ", context_str)
             prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
+            logger.debug(f"prompt: {prompt}")
         else:
             prompt = query
-        logger.debug(f"prompt: {prompt}")
         self.history.append([prompt, ''])
         response = ""
         for new_text in self.stream_generate_answer(
@@ -504,8 +517,29 @@ class Rag:
         self.history[-1][1] = response
         return response, reference_results
 
-    def query(self, query: str, **kwargs):
-        return self.predict(query, **kwargs)
+    def save_corpus_text(self):
+        if not self.corpus_files:
+            logger.warning("No hay archivos de corpus para guardar.")
+            return
+
+        corpus_text_file = os.path.join("corpus_embs/", "corpus_text.txt")
+
+        with open(corpus_text_file, 'w', encoding='utf-8') as f:
+            for chunk in self.sim_model.corpus.values():
+                f.write(chunk + "\n\n")  # Añade dos saltos de línea entre chunks para mejor legibilidad
+
+        logger.info(f"Texto del corpus guardado en: {corpus_text_file}")
+        return corpus_text_file
+
+    def load_corpus_text(self, emb_dir: str):
+        corpus_text_file = os.path.join("corpus_embs/", "corpus_text.txt")
+        if os.path.exists(corpus_text_file):
+            with open(corpus_text_file, 'r', encoding='utf-8') as f:
+                corpus_text = f.read().split("\n\n")  # Asumiendo que usamos dos saltos de línea como separador
+            self.sim_model.corpus = {i: chunk.strip() for i, chunk in enumerate(corpus_text) if chunk.strip()}
+            logger.info(f"Texto del corpus cargado desde: {corpus_text_file}")
+        else:
+            logger.warning(f"No se encontró el archivo de texto del corpus en: {corpus_text_file}")
 
     def save_corpus_emb(self):
         dir_name = self.get_file_hash(self.corpus_files)
@@ -517,13 +551,15 @@ class Rag:
 
     def load_corpus_emb(self, emb_dir: str):
         if hasattr(self.sim_model, 'load_corpus_embeddings'):
-            logger.debug(f"Loading corpus embeddings from {emb_dir}")
+            logger.debug(f"Cargando incrustaciones del corpus desde {emb_dir}")
             self.sim_model.load_corpus_embeddings(emb_dir)
+            # Cargar el texto del corpus
+            self.load_corpus_text(emb_dir)
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--sim_model_name", type=str, default="hiiamsid/sentence_similarity_spanish_es")
+    parser.add_argument("--sim_model_name", type=str, default="sentence-transformers/all-mpnet-base-v2")
     parser.add_argument("--gen_model_type", type=str, default="auto")
     parser.add_argument("--gen_model_name", type=str, default="Qwen/Qwen2-0.5B-Instruct")
     parser.add_argument("--lora_model", type=str, default=None)
@@ -538,7 +574,7 @@ if __name__ == "__main__":
     args = parser.parse_args()
     print(args)
     sim_model = BertSimilarity(model_name_or_path=args.sim_model_name, device=args.device)
-    m = Rag(
+    m = ChatPDF(
         similarity_model=sim_model,
         generate_model_type=args.gen_model_type,
         generate_model_name_or_path=args.gen_model_name,
@@ -551,4 +587,30 @@ if __name__ == "__main__":
         corpus_files=args.corpus_files.split(','),
         num_expand_context_chunk=args.num_expand_context_chunk,
         rerank_model_name_or_path=args.rerank_model_name,
-    )
+    )
+
+    # Comprobar si existen incrustaciones guardadas
+    dir_name = m.get_file_hash(args.corpus_files.split(','))
+    save_dir = os.path.join(m.save_corpus_emb_dir, dir_name)
+
+    if os.path.exists(save_dir):
+        # Cargar las incrustaciones guardadas
+        m.load_corpus_emb(save_dir)
+        print(f"Incrustaciones del corpus cargadas desde: {save_dir}")
+    else:
+        # Procesar el corpus y guardar las incrustaciones
+        m.add_corpus(args.corpus_files.split(','))
+        save_dir = m.save_corpus_emb()
+        # Guardar el texto del corpus
+        m.save_corpus_text()
+        print(f"Las incrustaciones del corpus se han guardado en: {save_dir}")
+
+    while True:
+        query = input("\nEnter a query: ")
+        if query == "exit":
+            break
+        if query.strip() == "":
+            continue
+        r, refs = m.predict(query)
+        print(r, refs)
+        print("\nRespuesta: ", r)