Update chatpdf.py

chatpdf.py

@@ -1,3 +1,8 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: 
+"""
 import argparse
 import hashlib
 import os
@@ -38,24 +43,6 @@ MODEL_CLASSES = {
     "auto": (AutoModelForCausalLM, AutoTokenizer),
 }
 
-PROMPT_TEMPLATE1 = """基于以下已知信息，简洁和专业的来回答用户的问题。
-如果无法从中得到答案，请说 "根据已知信息无法回答该问题" 或 "没有提供足够的相关信息"，不允许在答案中添加编造成分，答案请使用中文。
-
-已知内容:
-{context_str}
-
-问题:
-{query_str}
-"""
-PROMPT_TEMPLATE1 = """Utiliza la siguiente información para responder a la pregunta del usuario.
-Si no sabes la respuesta, di simplemente que no la sabes, no intentes inventarte una respuesta.
-
-Contexto: {context_str}
-Pregunta: {query_str}
-
-Devuelve sólo la respuesta útil que aparece a continuación y nada más, y ésta debe estar en Español.
-Respuesta útil:
-"""
 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 
@@ -69,6 +56,7 @@ Pregunta:
 """
 
 
+
 class SentenceSplitter:
     def __init__(self, chunk_size: int = 250, chunk_overlap: int = 50):
         self.chunk_size = chunk_size
@@ -134,7 +122,7 @@ class SentenceSplitter:
 
 
 
-class ChatPDF:
+class Rag:
     def __init__(
             self,
             similarity_model: SimilarityABC = None,
@@ -151,8 +139,8 @@ class ChatPDF:
             rerank_model_name_or_path: str = None,
             enable_history: bool = False,
             num_expand_context_chunk: int = 2,
-            similarity_top_k: int = 5,
-            rerank_top_k: int =3,
+            similarity_top_k: int = 10,
+            rerank_top_k: int = 3,
     ):
         """
         Init RAG model.
@@ -188,7 +176,7 @@ class ChatPDF:
         if similarity_model is not None:
             self.sim_model = similarity_model
         else:
-            m1 = BertSimilarity(model_name_or_path="jaimevera1107/all-MiniLM-L6-v2-similarity-es", device=self.device)
+            m1 = BertSimilarity(model_name_or_path="shibing624/text2vec-base-multilingual", device=self.device)
             m2 = BM25Similarity()
             default_sim_model = EnsembleSimilarity(similarities=[m1, m2], weights=[0.5, 0.5], c=2)
             self.sim_model = default_sim_model
@@ -205,7 +193,7 @@ class ChatPDF:
             self.add_corpus(corpus_files)
         self.save_corpus_emb_dir = save_corpus_emb_dir
         if rerank_model_name_or_path is None:
-            rerank_model_name_or_path = "BAAI/bge-reranker-large"
+            rerank_model_name_or_path = "BAAI/bge-reranker-base"
         if rerank_model_name_or_path:
             self.rerank_tokenizer = AutoTokenizer.from_pretrained(rerank_model_name_or_path)
             self.rerank_model = AutoModelForSequenceClassification.from_pretrained(rerank_model_name_or_path)
@@ -255,14 +243,14 @@ class ChatPDF:
         try:
             model.generation_config = GenerationConfig.from_pretrained(gen_model_name_or_path, trust_remote_code=True)
         except Exception as e:
-            logger.warning(f"No se pudo cargar la configuración de generación desde {gen_model_name_or_path}, {e}")
+            logger.warning(f"Failed to load generation config from {gen_model_name_or_path}, {e}")
         if peft_name:
             model = PeftModel.from_pretrained(
                 model,
                 peft_name,
                 torch_dtype="auto",
             )
-            logger.info(f"Modelo peft cargado desde {peft_name}")
+            logger.info(f"Loaded peft model from {peft_name}")
         model.eval()
         return model, tokenizer
 
@@ -353,9 +341,6 @@ class ChatPDF:
                 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
                     if text[-1] in ['.', '!', '?', '。', '！', '？', '…', ';', '；', ':', '：', '”', '’', '）', '】', '》', '」',
                                     '』', '〕', '〉', '》', '〗', '〞', '〟', '»', '"', "'", ')', ']', '}']:
@@ -422,9 +407,10 @@ class ChatPDF:
         sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
         # Get reference results from corpus
         hit_chunk_dict = dict()
-        for query_id, id_score_dict in sim_contents.items():
-            for corpus_id, s in id_score_dict.items():
-                hit_chunk = self.sim_model.corpus[corpus_id]
+        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
 
@@ -462,16 +448,15 @@ class ChatPDF:
             self.history = []
         if self.sim_model.corpus:
             reference_results = self.get_reference_results(query)
-            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)))
+            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 = ''
             prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
-            logger.debug(f"prompt: {prompt}")
         else:
             prompt = query
-            logger.debug(prompt)
+        logger.debug(f"prompt: {prompt}")
         self.history.append([prompt, ''])
         response = ""
         for new_text in self.stream_generate_answer(
@@ -496,20 +481,15 @@ class ChatPDF:
             self.history = []
         if self.sim_model.corpus:
             reference_results = self.get_reference_results(query)
-
-            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)
+            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 = ''
             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(
@@ -522,29 +502,8 @@ class ChatPDF:
         self.history[-1][1] = response
         return response, reference_results
 
-    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 query(self, query: str, **kwargs):
+        return self.predict(query, **kwargs)
 
     def save_corpus_emb(self):
         dir_name = self.get_file_hash(self.corpus_files)
@@ -556,20 +515,18 @@ class ChatPDF:
 
     def load_corpus_emb(self, emb_dir: str):
         if hasattr(self.sim_model, 'load_corpus_embeddings'):
-            logger.debug(f"Cargando incrustaciones del corpus desde {emb_dir}")
+            logger.debug(f"Loading corpus embeddings from {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="jaimevera1107/all-MiniLM-L6-v2-similarity-es")
+    parser.add_argument("--sim_model_name", type=str, default="shibing624/text2vec-base-multilingual")
     parser.add_argument("--gen_model_type", type=str, default="auto")
     parser.add_argument("--gen_model_name", type=str, default="LenguajeNaturalAI/leniachat-qwen2-1.5B-v0")
     parser.add_argument("--lora_model", type=str, default=None)
     parser.add_argument("--rerank_model_name", type=str, default="")
-    parser.add_argument("--corpus_files", type=str, default="Acuerdo009.pdf")
+    parser.add_argument("--corpus_files", type=str, default="data/sample.pdf")
     parser.add_argument("--device", type=str, default=None)
     parser.add_argument("--int4", action='store_true', help="use int4 quantization")
     parser.add_argument("--int8", action='store_true', help="use int8 quantization")
@@ -579,7 +536,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 = ChatPDF(
+    m = Rag(
         similarity_model=sim_model,
         generate_model_type=args.gen_model_type,
         generate_model_name_or_path=args.gen_model_name,
@@ -592,30 +549,4 @@ 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)
+    )