Update chatpdf.py

chatpdf.py

@@ -1,3 +1,8 @@
+# -*- coding: utf-8 -*-
+"""
+@author:XuMing([email protected])
+@description: 
+"""
 import argparse
 import hashlib
 import os
@@ -13,7 +18,6 @@ from similarities import (
     EnsembleSimilarity,
     BertSimilarity,
     BM25Similarity,
-    TfidfSimilarity
 )
 from similarities.similarity import SimilarityABC
 from transformers import (
@@ -50,6 +54,7 @@ Pregunta:
 {query_str}
 """
 
+
 class SentenceSplitter:
     def __init__(self, chunk_size: int = 250, chunk_overlap: int = 50):
         self.chunk_size = chunk_size
@@ -62,7 +67,7 @@ class SentenceSplitter:
             return self._split_english_text(text)
 
     def _split_chinese_text(self, text: str) -> List[str]:
-        sentence_endings = {'\n', '。', '！', '？', '；', '…'}  # puntuación al final de una frase
+        sentence_endings = {'\n', '。', '！', '？', '；', '…'}  # 句末标点符号
         chunks, current_chunk = [], ''
         for word in jieba.cut(text):
             if len(current_chunk) + len(word) > self.chunk_size:
@@ -80,16 +85,22 @@ class SentenceSplitter:
         return chunks
 
     def _split_english_text(self, text: str) -> List[str]:
-        # Dividir el texto inglés por frases utilizando expresiones regulares
+        # 使用正则表达式按句子分割英文文本
         sentences = re.split(r'(?<=[.!?])\s+', text.replace('\n', ' '))
-        chunks, current_chunk = [], ''
+        chunks = []
+        current_chunk = ''
         for sentence in sentences:
-            if len(current_chunk) + len(sentence) <= self.chunk_size or not current_chunk:
+            if len(current_chunk) + len(sentence) <= self.chunk_size:
                 current_chunk += (' ' if current_chunk else '') + sentence
             else:
-                chunks.append(current_chunk)
-                current_chunk = sentence
-        if current_chunk:  # Añade el último trozo
+                if len(sentence) > self.chunk_size:
+                    for i in range(0, len(sentence), self.chunk_size):
+                        chunks.append(sentence[i:i + self.chunk_size])
+                    current_chunk = ''
+                else:
+                    chunks.append(current_chunk)
+                    current_chunk = sentence
+        if current_chunk:  # Add the last chunk
             chunks.append(current_chunk)
 
         if self.chunk_overlap > 0 and len(chunks) > 1:
@@ -98,7 +109,7 @@ class SentenceSplitter:
         return chunks
 
     def _is_has_chinese(self, text: str) -> bool:
-        # comprobar si contiene caracteres chinos
+        # check if contains chinese characters
         if any("\u4e00" <= ch <= "\u9fff" for ch in text):
             return True
         else:
@@ -114,7 +125,7 @@ class SentenceSplitter:
         return overlapped_chunks
 
 
-class ChatPDF:
+class Rag:
     def __init__(
             self,
             similarity_model: SimilarityABC = None,
@@ -122,7 +133,7 @@ class ChatPDF:
             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/",
+            save_corpus_emb_dir: str = "./corpus_embs/",
             device: str = None,
             int8: bool = False,
             int4: bool = False,
@@ -131,8 +142,8 @@ class ChatPDF:
             rerank_model_name_or_path: str = None,
             enable_history: bool = False,
             num_expand_context_chunk: int = 2,
-            similarity_top_k: int = 15,
-            rerank_top_k: int = 5,
+            similarity_top_k: int = 10,
+            rerank_top_k: int = 3,
     ):
         """
         Init RAG model.
@@ -171,8 +182,7 @@ class ChatPDF:
             m1 = BertSimilarity(model_name_or_path="sentence-transformers/all-mpnet-base-v2", device=self.device)
             m2 = BM25Similarity()
             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
+            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,
@@ -237,14 +247,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
 
@@ -335,7 +345,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
@@ -408,12 +417,9 @@ class ChatPDF:
             # 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)
-
-        # Procesar los resultados de similitud
         reference_results = []
-
+        sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
+        # Get reference results from corpus
         hit_chunk_dict = dict()
         threshold_score = 0.5  # Establece un umbral para filtrar fragmentos irrelevantes
 
@@ -423,6 +429,7 @@ class ChatPDF:
                     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
@@ -447,9 +454,9 @@ class ChatPDF:
     def predict_stream(
             self,
             query: str,
-            max_length: int = 256,
-            context_len: int = 1024,
-            temperature: float = 0.5,
+            max_length: int = 512,
+            context_len: int = 2048,
+            temperature: float = 0.7,
     ):
         """Generate predictions stream."""
         stop_str = self.tokenizer.eos_token if self.tokenizer.eos_token else "</s>"
@@ -457,16 +464,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)[:]
+            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(
@@ -481,9 +487,9 @@ class ChatPDF:
     def predict(
             self,
             query: str,
-            max_length: int = 256,
-            context_len: int = 1024,
-            temperature: float = 0.5,
+            max_length: int = 512,
+            context_len: int = 2048,
+            temperature: float = 0.7,
     ):
         """Query from corpus."""
         reference_results = []
@@ -491,20 +497,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)[:]
+            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(
@@ -517,29 +518,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)
@@ -551,10 +531,8 @@ 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__":
@@ -564,7 +542,7 @@ if __name__ == "__main__":
     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="")
-    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")
@@ -574,7 +552,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,
@@ -588,29 +566,5 @@ if __name__ == "__main__":
         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)
+    r, refs = m.predict('自然语言中的非平行迁移是指什么？')
+    print(r, refs)