Update chatpdf.py

chatpdf.py

@@ -16,8 +16,13 @@ from similarities import (
 )
 from similarities.similarity import SimilarityABC
 from transformers import (
+    AutoModel,
     AutoModelForCausalLM,
     AutoTokenizer,
+    BloomForCausalLM,
+    BloomTokenizerFast,
+    LlamaTokenizer,
+    LlamaForCausalLM,
     TextIteratorStreamer,
     GenerationConfig,
     AutoModelForSequenceClassification,
@@ -26,9 +31,22 @@ from transformers import (
 jieba.setLogLevel("ERROR")
 
 MODEL_CLASSES = {
+    "bloom": (BloomForCausalLM, BloomTokenizerFast),
+    "chatglm": (AutoModel, AutoTokenizer),
+    "llama": (LlamaForCausalLM, LlamaTokenizer),
+    "baichuan": (AutoModelForCausalLM, AutoTokenizer),
     "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.
 
@@ -41,7 +59,7 @@ 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 
-inventados en la respuesta.
+inventados en la respuesta, y ésta debe estar en Español.
 
 Contenido conocido:
 {context_str}
@@ -81,7 +99,7 @@ class SentenceSplitter:
         return chunks
 
     def _split_english_text(self, text: str) -> List[str]:
-        # División de texto inglés por frases mediante expresiones regulares
+        # Dividir el texto inglés por frases utilizando expresiones regulares
         sentences = re.split(r'(?<=[.!?])\s+', text.replace('\n', ' '))
         chunks, current_chunk = [], ''
         for sentence in sentences:
@@ -90,7 +108,7 @@ class SentenceSplitter:
             else:
                 chunks.append(current_chunk)
                 current_chunk = sentence
-        if current_chunk:  # Add the last chunk
+        if current_chunk:  # Añade el último trozo
             chunks.append(current_chunk)
 
         if self.chunk_overlap > 0 and len(chunks) > 1:
@@ -99,14 +117,14 @@ class SentenceSplitter:
         return chunks
 
     def _is_has_chinese(self, text: str) -> bool:
-        # check if contains chinese characters
+        # comprobar si contiene caracteres chinos
         if any("\u4e00" <= ch <= "\u9fff" for ch in text):
             return True
         else:
             return False
 
     def _handle_overlap(self, chunks: List[str]) -> List[str]:
-        # Tratamiento de los solapamientos entre bloques
+        # 处理块间重叠
         overlapped_chunks = []
         for i in range(len(chunks) - 1):
             chunk = chunks[i] + ' ' + chunks[i + 1][:self.chunk_overlap]
@@ -115,12 +133,13 @@ class SentenceSplitter:
         return overlapped_chunks
 
 
+
 class ChatPDF:
     def __init__(
             self,
             similarity_model: SimilarityABC = None,
             generate_model_type: str = "auto",
-            generate_model_name_or_path: str = "LenguajeNaturalAI/leniachat-qwen2-1.5B-v0",
+            generate_model_name_or_path: str = "Qwen/Qwen2-0.5B-Instruct",
             lora_model_name_or_path: str = None,
             corpus_files: Union[str, List[str]] = None,
             save_corpus_emb_dir: str = "corpus_embs/",
@@ -132,10 +151,28 @@ class ChatPDF:
             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 = 5,
+            rerank_top_k: int =3,
     ):
-
+        """
+        Init RAG model.
+        :param similarity_model: similarity model, default None, if set, will use it instead of EnsembleSimilarity
+        :param generate_model_type: generate model type
+        :param generate_model_name_or_path: generate model name or path
+        :param lora_model_name_or_path: lora model name or path
+        :param corpus_files: corpus files
+        :param save_corpus_emb_dir: save corpus embeddings dir, default ./corpus_embs/
+        :param device: device, default None, auto select gpu or cpu
+        :param int8: use int8 quantization, default False
+        :param int4: use int4 quantization, default False
+        :param chunk_size: chunk size, default 250
+        :param chunk_overlap: chunk overlap, default 0, can not set to > 0 if num_expand_context_chunk > 0
+        :param rerank_model_name_or_path: rerank model name or path, default 'BAAI/bge-reranker-base'
+        :param enable_history: enable history, default False
+        :param num_expand_context_chunk: num expand context chunk, default 2, if set to 0, will not expand context chunk
+        :param similarity_top_k: similarity_top_k, default 5, similarity model search k corpus chunks
+        :param rerank_top_k: rerank_top_k, default 3, rerank model search k corpus chunks
+        """
         if torch.cuda.is_available():
             default_device = torch.device(0)
         elif torch.backends.mps.is_available():
@@ -151,7 +188,7 @@ class ChatPDF:
         if similarity_model is not None:
             self.sim_model = similarity_model
         else:
-            m1 = BertSimilarity(model_name_or_path="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", 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
@@ -168,7 +205,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 = "maidalun1020/bce-reranker-base_v1"
+            rerank_model_name_or_path = "BAAI/bge-reranker-large"
         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)
@@ -252,7 +289,7 @@ class ChatPDF:
             repetition_penalty=1.0,
             context_len=2048
     ):
-        streamer = TextIteratorStreamer(self.tokenizer, timeout=520.0, skip_prompt=True, skip_special_tokens=True)
+        streamer = TextIteratorStreamer(self.tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
         input_ids = self._get_chat_input()
         max_src_len = context_len - max_new_tokens - 8
         input_ids = input_ids[-max_src_len:]
@@ -383,29 +420,14 @@ class ChatPDF:
         """
         reference_results = []
         sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
-    
-        # Verificar si sim_contents es una lista o un diccionario
-        if isinstance(sim_contents, list):
-            for item in sim_contents:
-                # Ajustar según la estructura real de item
-                corpus_id = item[0] if isinstance(item, (list, tuple)) else item  # Asegurarse de que corpus_id sea el valor correcto
-                if isinstance(corpus_id, dict):
-                    # Extraer el valor necesario si corpus_id es un diccionario
-                    corpus_id = next(iter(corpus_id.keys()))  # Tomar la primera clave como ejemplo
-                if corpus_id in self.sim_model.corpus:
-                    hit_chunk = self.sim_model.corpus[corpus_id]
-                    reference_results.append(hit_chunk)
-    
-        elif isinstance(sim_contents, dict):
-            for query_id, id_score_dict in sim_contents.items():
-                for corpus_id, s in id_score_dict.items():
-                    if corpus_id in self.sim_model.corpus:
-                        hit_chunk = self.sim_model.corpus[corpus_id]
-                        reference_results.append(hit_chunk)
-        else:
-            logger.error(f"Unexpected type for sim_contents: {type(sim_contents)}")
+        # 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]
+                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
@@ -444,7 +466,7 @@ class ChatPDF:
                 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("context_str: " , (context_len - len(PROMPT_TEMPLATE)))
+            print("gggggg: ", (context_len - len(PROMPT_TEMPLATE)))
             prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
             logger.debug(f"prompt: {prompt}")
         else:
@@ -478,12 +500,12 @@ class ChatPDF:
             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_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: ", (context_len - len(PROMPT_TEMPLATE)))
             print(".......................................................")
             context_str = '\n'.join(reference_results)[:]
-            #print("context_str: ", context_str)
+            print("context_str: ", context_str)
             prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
             logger.debug(f"prompt: {prompt}")
         else:
@@ -500,19 +522,6 @@ class ChatPDF:
         self.history[-1][1] = response
         return response, reference_results
 
-    def save_corpus_emb(self):
-        dir_name = self.get_file_hash(self.corpus_files)
-        save_dir = os.path.join(self.save_corpus_emb_dir, dir_name)
-        if hasattr(self.sim_model, 'save_corpus_embeddings'):
-            self.sim_model.save_corpus_embeddings(save_dir)
-            logger.debug(f"Saving corpus embeddings to {save_dir}")
-        return save_dir
-
-    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}")
-            self.sim_model.load_corpus_embeddings(emb_dir)
-
     def save_corpus_text(self):
         if not self.corpus_files:
             logger.warning("No hay archivos de corpus para guardar.")
@@ -537,20 +546,36 @@ class ChatPDF:
         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)
+        save_dir = os.path.join(self.save_corpus_emb_dir, dir_name)
+        if hasattr(self.sim_model, 'save_corpus_embeddings'):
+            self.sim_model.save_corpus_embeddings(save_dir)
+            logger.debug(f"Saving corpus embeddings to {save_dir}")
+        return save_dir
+
+    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}")
+            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="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    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("--gen_model_name", type=str, default="Qwen/Qwen2-0.5B-Instruct")
     parser.add_argument("--lora_model", type=str, default=None)
-    parser.add_argument("--rerank_model_name", type=str, default="maidalun1020/bce-reranker-base_v1")
-    parser.add_argument("--corpus_files", type=str, default="docs/corpus.txt")
+    parser.add_argument("--rerank_model_name", type=str, default="")
+    parser.add_argument("--corpus_files", type=str, default="Acuerdo009.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")
     parser.add_argument("--chunk_size", type=int, default=220)
-    parser.add_argument("--chunk_overlap", type=int, default=50)
-    parser.add_argument("--num_expand_context_chunk", type=int, default=2)
+    parser.add_argument("--chunk_overlap", type=int, default=0)
+    parser.add_argument("--num_expand_context_chunk", type=int, default=1)
     args = parser.parse_args()
     print(args)
     sim_model = BertSimilarity(model_name_or_path=args.sim_model_name, device=args.device)
@@ -568,7 +593,6 @@ if __name__ == "__main__":
         num_expand_context_chunk=args.num_expand_context_chunk,
         rerank_model_name_or_path=args.rerank_model_name,
     )
-    logger.info(f"chatpdf model: {m}")
 
     # Comprobar si existen incrustaciones guardadas
     dir_name = m.get_file_hash(args.corpus_files.split(','))
@@ -594,4 +618,4 @@ if __name__ == "__main__":
             continue
         r, refs = m.predict(query)
         print(r, refs)
-        print("\nRespuesta: ", r)
+        print("\nRespuesta: ", r)