Update chatpdf.py

chatpdf.py

@@ -1,582 +1,591 @@
-import argparse
-import hashlib
-import os
-import re
-from threading import Thread
-from typing import Union, List
-
-import jieba
-import torch
-from loguru import logger
-from peft import PeftModel
-from similarities import (
-    EnsembleSimilarity,
-    BertSimilarity,
-    BM25Similarity,
-)
-from similarities.similarity import SimilarityABC
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-    GenerationConfig,
-    AutoModelForSequenceClassification,
-)
-
-jieba.setLogLevel("ERROR")
-
-MODEL_CLASSES = {
-    "auto": (AutoModelForCausalLM, AutoTokenizer),
-}
-
-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 
-inventados en la respuesta.
-
-Contenido conocido:
-{context_str}
-
-Pregunta:
-{query_str}
-"""
-
-
-class SentenceSplitter:
-    def __init__(self, chunk_size: int = 250, chunk_overlap: int = 50):
-        self.chunk_size = chunk_size
-        self.chunk_overlap = chunk_overlap
-
-    def split_text(self, text: str) -> List[str]:
-        if self._is_has_chinese(text):
-            return self._split_chinese_text(text)
-        else:
-            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
-        chunks, current_chunk = [], ''
-        for word in jieba.cut(text):
-            if len(current_chunk) + len(word) > self.chunk_size:
-                chunks.append(current_chunk.strip())
-                current_chunk = word
-            else:
-                current_chunk += word
-            if word[-1] in sentence_endings and len(current_chunk) > self.chunk_size - self.chunk_overlap:
-                chunks.append(current_chunk.strip())
-                current_chunk = ''
-        if current_chunk:
-            chunks.append(current_chunk.strip())
-        if self.chunk_overlap > 0 and len(chunks) > 1:
-            chunks = self._handle_overlap(chunks)
-        return chunks
-
-    def _split_english_text(self, text: str) -> List[str]:
-        # División de texto inglés por frases mediante expresiones regulares
-        sentences = re.split(r'(?<=[.!?])\s+', text.replace('\n', ' '))
-        chunks, current_chunk = [], ''
-        for sentence in sentences:
-            if len(current_chunk) + len(sentence) <= self.chunk_size or not current_chunk:
-                current_chunk += (' ' if current_chunk else '') + sentence
-            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:
-            chunks = self._handle_overlap(chunks)
-
-        return chunks
-
-    def _is_has_chinese(self, text: str) -> bool:
-        # check if contains chinese characters
-        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]
-            overlapped_chunks.append(chunk.strip())
-        overlapped_chunks.append(chunks[-1])
-        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",
-            lora_model_name_or_path: str = None,
-            corpus_files: Union[str, List[str]] = None,
-            save_corpus_emb_dir: str = "corpus_embs/",
-            device: str = None,
-            int8: bool = False,
-            int4: bool = False,
-            chunk_size: int = 250,
-            chunk_overlap: int = 0,
-            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
-    ):
-
-        if torch.cuda.is_available():
-            default_device = torch.device(0)
-        elif torch.backends.mps.is_available():
-            default_device = torch.device('cpu')
-        else:
-            default_device = torch.device('cpu')
-        self.device = device or default_device
-        if num_expand_context_chunk > 0 and chunk_overlap > 0:
-            logger.warning(f" 'num_expand_context_chunk' and 'chunk_overlap' cannot both be greater than zero. "
-                           f" 'chunk_overlap' has been set to zero by default.")
-            chunk_overlap = 0
-        self.text_splitter = SentenceSplitter(chunk_size, chunk_overlap)
-        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)
-            m2 = BM25Similarity()
-            default_sim_model = EnsembleSimilarity(similarities=[m1, m2], weights=[0.5, 0.5], c=2)
-            self.sim_model = default_sim_model
-        self.gen_model, self.tokenizer = self._init_gen_model(
-            generate_model_type,
-            generate_model_name_or_path,
-            peft_name=lora_model_name_or_path,
-            int8=int8,
-            int4=int4,
-        )
-        self.history = []
-        self.corpus_files = corpus_files
-        if corpus_files:
-            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"
-        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)
-            self.rerank_model.to(self.device)
-            self.rerank_model.eval()
-        else:
-            self.rerank_model = None
-            self.rerank_tokenizer = None
-        self.enable_history = enable_history
-        self.similarity_top_k = similarity_top_k
-        self.num_expand_context_chunk = num_expand_context_chunk
-        self.rerank_top_k = rerank_top_k
-
-    def __str__(self):
-        return f"Similarity model: {self.sim_model}, Generate model: {self.gen_model}"
-
-    def _init_gen_model(
-            self,
-            gen_model_type: str,
-            gen_model_name_or_path: str,
-            peft_name: str = None,
-            int8: bool = False,
-            int4: bool = False,
-    ):
-        """Init generate model."""
-        if int8 or int4:
-            device_map = None
-        else:
-            device_map = "auto"
-        model_class, tokenizer_class = MODEL_CLASSES[gen_model_type]
-        tokenizer = tokenizer_class.from_pretrained(gen_model_name_or_path, trust_remote_code=True)
-        model = model_class.from_pretrained(
-            gen_model_name_or_path,
-            load_in_8bit=int8 if gen_model_type not in ['baichuan', 'chatglm'] else False,
-            load_in_4bit=int4 if gen_model_type not in ['baichuan', 'chatglm'] else False,
-            torch_dtype="auto",
-            device_map=device_map,
-            trust_remote_code=True,
-        )
-        if self.device == torch.device('cpu'):
-            model.float()
-        if gen_model_type in ['baichuan', 'chatglm']:
-            if int4:
-                model = model.quantize(4).cuda()
-            elif int8:
-                model = model.quantize(8).cuda()
-        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}")
-        if peft_name:
-            model = PeftModel.from_pretrained(
-                model,
-                peft_name,
-                torch_dtype="auto",
-            )
-            logger.info(f"Modelo peft cargado desde {peft_name}")
-        model.eval()
-        return model, tokenizer
-
-    def _get_chat_input(self):
-        messages = []
-        for conv in self.history:
-            if conv and len(conv) > 0 and conv[0]:
-                messages.append({'role': 'user', 'content': conv[0]})
-            if conv and len(conv) > 1 and conv[1]:
-                messages.append({'role': 'assistant', 'content': conv[1]})
-        input_ids = self.tokenizer.apply_chat_template(
-            conversation=messages,
-            tokenize=True,
-            add_generation_prompt=True,
-            return_tensors='pt'
-        )
-        return input_ids.to(self.gen_model.device)
-
-    @torch.inference_mode()
-    def stream_generate_answer(
-            self,
-            max_new_tokens=512,
-            temperature=0.7,
-            repetition_penalty=1.0,
-            context_len=2048
-    ):
-        streamer = TextIteratorStreamer(self.tokenizer, timeout=520.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:]
-        generation_kwargs = dict(
-            input_ids=input_ids,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            do_sample=True,
-            repetition_penalty=repetition_penalty,
-            streamer=streamer,
-        )
-        thread = Thread(target=self.gen_model.generate, kwargs=generation_kwargs)
-        thread.start()
-
-        yield from streamer
-
-    def add_corpus(self, files: Union[str, List[str]]):
-        """Load document files."""
-        if isinstance(files, str):
-            files = [files]
-        for doc_file in files:
-            if doc_file.endswith('.pdf'):
-                corpus = self.extract_text_from_pdf(doc_file)
-            elif doc_file.endswith('.docx'):
-                corpus = self.extract_text_from_docx(doc_file)
-            elif doc_file.endswith('.md'):
-                corpus = self.extract_text_from_markdown(doc_file)
-            else:
-                corpus = self.extract_text_from_txt(doc_file)
-            full_text = '\n'.join(corpus)
-            chunks = self.text_splitter.split_text(full_text)
-            self.sim_model.add_corpus(chunks)
-        self.corpus_files = files
-        logger.debug(f"files: {files}, corpus size: {len(self.sim_model.corpus)}, top3: "
-                     f"{list(self.sim_model.corpus.values())[:3]}")
-
-    @staticmethod
-    def get_file_hash(fpaths):
-        hasher = hashlib.md5()
-        target_file_data = bytes()
-        if isinstance(fpaths, str):
-            fpaths = [fpaths]
-        for fpath in fpaths:
-            with open(fpath, 'rb') as file:
-                chunk = file.read(1024 * 1024)  # read only first 1MB
-                hasher.update(chunk)
-                target_file_data += chunk
-
-        hash_name = hasher.hexdigest()[:32]
-        return hash_name
-
-    @staticmethod
-    def extract_text_from_pdf(file_path: str):
-        """Extract text content from a PDF file."""
-        import PyPDF2
-        contents = []
-        with open(file_path, 'rb') as f:
-            pdf_reader = PyPDF2.PdfReader(f)
-            for page in pdf_reader.pages:
-                page_text = page.extract_text().strip()
-                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 ['.', '!', '?', '。', '！', '？', '…', ';', '；', ':', '：', '”', '’', '）', '】', '》', '」',
-                                    '』', '〕', '〉', '》', '〗', '〞', '〟', '»', '"', "'", ')', ']', '}']:
-                        contents.append(new_text)
-                        new_text = ''
-                if new_text:
-                    contents.append(new_text)
-        return contents
-
-    @staticmethod
-    def extract_text_from_txt(file_path: str):
-        """Extract text content from a TXT file."""
-        with open(file_path, 'r', encoding='utf-8') as f:
-            contents = [text.strip() for text in f.readlines() if text.strip()]
-        return contents
-
-    @staticmethod
-    def extract_text_from_docx(file_path: str):
-        """Extract text content from a DOCX file."""
-        import docx
-        document = docx.Document(file_path)
-        contents = [paragraph.text.strip() for paragraph in document.paragraphs if paragraph.text.strip()]
-        return contents
-
-    @staticmethod
-    def extract_text_from_markdown(file_path: str):
-        """Extract text content from a Markdown file."""
-        import markdown
-        from bs4 import BeautifulSoup
-        with open(file_path, 'r', encoding='utf-8') as f:
-            markdown_text = f.read()
-        html = markdown.markdown(markdown_text)
-        soup = BeautifulSoup(html, 'html.parser')
-        contents = [text.strip() for text in soup.get_text().splitlines() if text.strip()]
-        return contents
-
-    @staticmethod
-    def _add_source_numbers(lst):
-        """Add source numbers to a list of strings."""
-        return [f'[{idx + 1}]\t "{item}"' for idx, item in enumerate(lst)]
-
-    def _get_reranker_score(self, query: str, reference_results: List[str]):
-        """Get reranker score."""
-        pairs = []
-        for reference in reference_results:
-            pairs.append([query, reference])
-        with torch.no_grad():
-            inputs = self.rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
-            inputs_on_device = {k: v.to(self.rerank_model.device) for k, v in inputs.items()}
-            scores = self.rerank_model(**inputs_on_device, return_dict=True).logits.view(-1, ).float()
-
-        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 = []
-        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]
-                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
-                rerank_scores = self._get_reranker_score(query, reference_results)
-                logger.debug(f"rerank_scores: {rerank_scores}")
-                # Get rerank top k chunks
-                reference_results = [reference for reference, score in sorted(
-                    zip(reference_results, rerank_scores), key=lambda x: x[1], reverse=True)][:self.rerank_top_k]
-                hit_chunk_dict = {corpus_id: hit_chunk for corpus_id, hit_chunk in hit_chunk_dict.items() if
-                                  hit_chunk in reference_results}
-            # Expand reference context chunk
-            if self.num_expand_context_chunk > 0:
-                new_reference_results = []
-                for corpus_id, hit_chunk in hit_chunk_dict.items():
-                    expanded_reference = self.sim_model.corpus.get(corpus_id - 1, '') + hit_chunk
-                    for i in range(self.num_expand_context_chunk):
-                        expanded_reference += self.sim_model.corpus.get(corpus_id + i + 1, '')
-                    new_reference_results.append(expanded_reference)
-                reference_results = new_reference_results
-        return reference_results
-
-    def predict_stream(
-            self,
-            query: str,
-            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>"
-        if not self.enable_history:
-            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("context_str: " , (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(prompt)
-        self.history.append([prompt, ''])
-        response = ""
-        for new_text in self.stream_generate_answer(
-                max_new_tokens=max_length,
-                temperature=temperature,
-                context_len=context_len,
-        ):
-            if new_text != stop_str:
-                response += new_text
-                yield response
-
-    def predict(
-            self,
-            query: str,
-            max_length: int = 512,
-            context_len: int = 2048,
-            temperature: float = 0.7,
-    ):
-        """Query from corpus."""
-        reference_results = []
-        if not self.enable_history:
-            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)
-            prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
-            logger.debug(f"prompt: {prompt}")
-        else:
-            prompt = query
-        self.history.append([prompt, ''])
-        response = ""
-        for new_text in self.stream_generate_answer(
-                max_new_tokens=max_length,
-                temperature=temperature,
-                context_len=context_len,
-        ):
-            response += new_text
-        response = response.strip()
-        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.")
-            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}")
-
-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("--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="maidalun1020/bce-reranker-base_v1")
-    parser.add_argument("--corpus_files", type=str, default="docs/corpus.txt")
-    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)
-    args = parser.parse_args()
-    print(args)
-    sim_model = BertSimilarity(model_name_or_path=args.sim_model_name, device=args.device)
-    m = ChatPDF(
-        similarity_model=sim_model,
-        generate_model_type=args.gen_model_type,
-        generate_model_name_or_path=args.gen_model_name,
-        lora_model_name_or_path=args.lora_model,
-        device=args.device,
-        int4=args.int4,
-        int8=args.int8,
-        chunk_size=args.chunk_size,
-        chunk_overlap=args.chunk_overlap,
-        corpus_files=args.corpus_files.split(','),
-        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(','))
-    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)
+import argparse
+import hashlib
+import os
+import re
+from threading import Thread
+from typing import Union, List
+
+import jieba
+import torch
+from loguru import logger
+from peft import PeftModel
+from similarities import (
+    EnsembleSimilarity,
+    BertSimilarity,
+    BM25Similarity,
+)
+from similarities.similarity import SimilarityABC
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    TextIteratorStreamer,
+    GenerationConfig,
+    AutoModelForSequenceClassification,
+)
+
+jieba.setLogLevel("ERROR")
+
+MODEL_CLASSES = {
+    "auto": (AutoModelForCausalLM, AutoTokenizer),
+}
+
+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 
+inventados en la respuesta.
+
+Contenido conocido:
+{context_str}
+
+Pregunta:
+{query_str}
+"""
+
+
+class SentenceSplitter:
+    def __init__(self, chunk_size: int = 250, chunk_overlap: int = 50):
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+
+    def split_text(self, text: str) -> List[str]:
+        if self._is_has_chinese(text):
+            return self._split_chinese_text(text)
+        else:
+            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
+        chunks, current_chunk = [], ''
+        for word in jieba.cut(text):
+            if len(current_chunk) + len(word) > self.chunk_size:
+                chunks.append(current_chunk.strip())
+                current_chunk = word
+            else:
+                current_chunk += word
+            if word[-1] in sentence_endings and len(current_chunk) > self.chunk_size - self.chunk_overlap:
+                chunks.append(current_chunk.strip())
+                current_chunk = ''
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+        if self.chunk_overlap > 0 and len(chunks) > 1:
+            chunks = self._handle_overlap(chunks)
+        return chunks
+
+    def _split_english_text(self, text: str) -> List[str]:
+        # División de texto inglés por frases mediante expresiones regulares
+        sentences = re.split(r'(?<=[.!?])\s+', text.replace('\n', ' '))
+        chunks, current_chunk = [], ''
+        for sentence in sentences:
+            if len(current_chunk) + len(sentence) <= self.chunk_size or not current_chunk:
+                current_chunk += (' ' if current_chunk else '') + sentence
+            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:
+            chunks = self._handle_overlap(chunks)
+
+        return chunks
+
+    def _is_has_chinese(self, text: str) -> bool:
+        # check if contains chinese characters
+        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]
+            overlapped_chunks.append(chunk.strip())
+        overlapped_chunks.append(chunks[-1])
+        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",
+            lora_model_name_or_path: str = None,
+            corpus_files: Union[str, List[str]] = None,
+            save_corpus_emb_dir: str = "corpus_embs/",
+            device: str = None,
+            int8: bool = False,
+            int4: bool = False,
+            chunk_size: int = 250,
+            chunk_overlap: int = 0,
+            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
+    ):
+
+        if torch.cuda.is_available():
+            default_device = torch.device(0)
+        elif torch.backends.mps.is_available():
+            default_device = torch.device('cpu')
+        else:
+            default_device = torch.device('cpu')
+        self.device = device or default_device
+        if num_expand_context_chunk > 0 and chunk_overlap > 0:
+            logger.warning(f" 'num_expand_context_chunk' and 'chunk_overlap' cannot both be greater than zero. "
+                           f" 'chunk_overlap' has been set to zero by default.")
+            chunk_overlap = 0
+        self.text_splitter = SentenceSplitter(chunk_size, chunk_overlap)
+        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)
+            m2 = BM25Similarity()
+            default_sim_model = EnsembleSimilarity(similarities=[m1, m2], weights=[0.5, 0.5], c=2)
+            self.sim_model = default_sim_model
+        self.gen_model, self.tokenizer = self._init_gen_model(
+            generate_model_type,
+            generate_model_name_or_path,
+            peft_name=lora_model_name_or_path,
+            int8=int8,
+            int4=int4,
+        )
+        self.history = []
+        self.corpus_files = corpus_files
+        if corpus_files:
+            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"
+        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)
+            self.rerank_model.to(self.device)
+            self.rerank_model.eval()
+        else:
+            self.rerank_model = None
+            self.rerank_tokenizer = None
+        self.enable_history = enable_history
+        self.similarity_top_k = similarity_top_k
+        self.num_expand_context_chunk = num_expand_context_chunk
+        self.rerank_top_k = rerank_top_k
+
+    def __str__(self):
+        return f"Similarity model: {self.sim_model}, Generate model: {self.gen_model}"
+
+    def _init_gen_model(
+            self,
+            gen_model_type: str,
+            gen_model_name_or_path: str,
+            peft_name: str = None,
+            int8: bool = False,
+            int4: bool = False,
+    ):
+        """Init generate model."""
+        if int8 or int4:
+            device_map = None
+        else:
+            device_map = "auto"
+        model_class, tokenizer_class = MODEL_CLASSES[gen_model_type]
+        tokenizer = tokenizer_class.from_pretrained(gen_model_name_or_path, trust_remote_code=True)
+        model = model_class.from_pretrained(
+            gen_model_name_or_path,
+            load_in_8bit=int8 if gen_model_type not in ['baichuan', 'chatglm'] else False,
+            load_in_4bit=int4 if gen_model_type not in ['baichuan', 'chatglm'] else False,
+            torch_dtype="auto",
+            device_map=device_map,
+            trust_remote_code=True,
+        )
+        if self.device == torch.device('cpu'):
+            model.float()
+        if gen_model_type in ['baichuan', 'chatglm']:
+            if int4:
+                model = model.quantize(4).cuda()
+            elif int8:
+                model = model.quantize(8).cuda()
+        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}")
+        if peft_name:
+            model = PeftModel.from_pretrained(
+                model,
+                peft_name,
+                torch_dtype="auto",
+            )
+            logger.info(f"Modelo peft cargado desde {peft_name}")
+        model.eval()
+        return model, tokenizer
+
+    def _get_chat_input(self):
+        messages = []
+        for conv in self.history:
+            if conv and len(conv) > 0 and conv[0]:
+                messages.append({'role': 'user', 'content': conv[0]})
+            if conv and len(conv) > 1 and conv[1]:
+                messages.append({'role': 'assistant', 'content': conv[1]})
+        input_ids = self.tokenizer.apply_chat_template(
+            conversation=messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors='pt'
+        )
+        return input_ids.to(self.gen_model.device)
+
+    @torch.inference_mode()
+    def stream_generate_answer(
+            self,
+            max_new_tokens=512,
+            temperature=0.7,
+            repetition_penalty=1.0,
+            context_len=2048
+    ):
+        streamer = TextIteratorStreamer(self.tokenizer, timeout=520.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:]
+        generation_kwargs = dict(
+            input_ids=input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            repetition_penalty=repetition_penalty,
+            streamer=streamer,
+        )
+        thread = Thread(target=self.gen_model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        yield from streamer
+
+    def add_corpus(self, files: Union[str, List[str]]):
+        """Load document files."""
+        if isinstance(files, str):
+            files = [files]
+        for doc_file in files:
+            if doc_file.endswith('.pdf'):
+                corpus = self.extract_text_from_pdf(doc_file)
+            elif doc_file.endswith('.docx'):
+                corpus = self.extract_text_from_docx(doc_file)
+            elif doc_file.endswith('.md'):
+                corpus = self.extract_text_from_markdown(doc_file)
+            else:
+                corpus = self.extract_text_from_txt(doc_file)
+            full_text = '\n'.join(corpus)
+            chunks = self.text_splitter.split_text(full_text)
+            self.sim_model.add_corpus(chunks)
+        self.corpus_files = files
+        logger.debug(f"files: {files}, corpus size: {len(self.sim_model.corpus)}, top3: "
+                     f"{list(self.sim_model.corpus.values())[:3]}")
+
+    @staticmethod
+    def get_file_hash(fpaths):
+        hasher = hashlib.md5()
+        target_file_data = bytes()
+        if isinstance(fpaths, str):
+            fpaths = [fpaths]
+        for fpath in fpaths:
+            with open(fpath, 'rb') as file:
+                chunk = file.read(1024 * 1024)  # read only first 1MB
+                hasher.update(chunk)
+                target_file_data += chunk
+
+        hash_name = hasher.hexdigest()[:32]
+        return hash_name
+
+    @staticmethod
+    def extract_text_from_pdf(file_path: str):
+        """Extract text content from a PDF file."""
+        import PyPDF2
+        contents = []
+        with open(file_path, 'rb') as f:
+            pdf_reader = PyPDF2.PdfReader(f)
+            for page in pdf_reader.pages:
+                page_text = page.extract_text().strip()
+                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 ['.', '!', '?', '。', '！', '？', '…', ';', '；', ':', '：', '”', '’', '）', '】', '》', '」',
+                                    '』', '〕', '〉', '》', '〗', '〞', '〟', '»', '"', "'", ')', ']', '}']:
+                        contents.append(new_text)
+                        new_text = ''
+                if new_text:
+                    contents.append(new_text)
+        return contents
+
+    @staticmethod
+    def extract_text_from_txt(file_path: str):
+        """Extract text content from a TXT file."""
+        with open(file_path, 'r', encoding='utf-8') as f:
+            contents = [text.strip() for text in f.readlines() if text.strip()]
+        return contents
+
+    @staticmethod
+    def extract_text_from_docx(file_path: str):
+        """Extract text content from a DOCX file."""
+        import docx
+        document = docx.Document(file_path)
+        contents = [paragraph.text.strip() for paragraph in document.paragraphs if paragraph.text.strip()]
+        return contents
+
+    @staticmethod
+    def extract_text_from_markdown(file_path: str):
+        """Extract text content from a Markdown file."""
+        import markdown
+        from bs4 import BeautifulSoup
+        with open(file_path, 'r', encoding='utf-8') as f:
+            markdown_text = f.read()
+        html = markdown.markdown(markdown_text)
+        soup = BeautifulSoup(html, 'html.parser')
+        contents = [text.strip() for text in soup.get_text().splitlines() if text.strip()]
+        return contents
+
+    @staticmethod
+    def _add_source_numbers(lst):
+        """Add source numbers to a list of strings."""
+        return [f'[{idx + 1}]\t "{item}"' for idx, item in enumerate(lst)]
+
+    def _get_reranker_score(self, query: str, reference_results: List[str]):
+        """Get reranker score."""
+        pairs = []
+        for reference in reference_results:
+            pairs.append([query, reference])
+        with torch.no_grad():
+            inputs = self.rerank_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
+            inputs_on_device = {k: v.to(self.rerank_model.device) for k, v in inputs.items()}
+            scores = self.rerank_model(**inputs_on_device, return_dict=True).logits.view(-1, ).float()
+
+        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 = []
+    sim_contents = self.sim_model.most_similar(query, topn=self.similarity_top_k)
+    
+    # Ajustar según el tipo de retorno de sim_contents
+    if isinstance(sim_contents, dict):  # Si es un diccionario
+        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)
+    elif isinstance(sim_contents, list):  # Si es una lista
+        for item in sim_contents:
+            # Ajusta esto dependiendo de la estructura de los elementos de la lista
+            # Ejemplo: si es una lista de (corpus_id, score) tuplas
+            corpus_id, _ = item
+            hit_chunk = self.sim_model.corpus[corpus_id]
+            reference_results.append(hit_chunk)
+    
+    # Resto del código...
+    if reference_results:
+        if self.rerank_model is not None:
+            # Rerank reference results
+            rerank_scores = self._get_reranker_score(query, reference_results)
+            logger.debug(f"rerank_scores: {rerank_scores}")
+            # Get rerank top k chunks
+            reference_results = [reference for reference, score in sorted(
+                zip(reference_results, rerank_scores), key=lambda x: x[1], reverse=True)][:self.rerank_top_k]
+            hit_chunk_dict = {corpus_id: hit_chunk for corpus_id, hit_chunk in hit_chunk_dict.items() if
+                              hit_chunk in reference_results}
+        # Expand reference context chunk
+        if self.num_expand_context_chunk > 0:
+            new_reference_results = []
+            for corpus_id, hit_chunk in hit_chunk_dict.items():
+                expanded_reference = self.sim_model.corpus.get(corpus_id - 1, '') + hit_chunk
+                for i in range(self.num_expand_context_chunk):
+                    expanded_reference += self.sim_model.corpus.get(corpus_id + i + 1, '')
+                new_reference_results.append(expanded_reference)
+            reference_results = new_reference_results
+    return reference_results
+
+
+    def predict_stream(
+            self,
+            query: str,
+            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>"
+        if not self.enable_history:
+            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("context_str: " , (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(prompt)
+        self.history.append([prompt, ''])
+        response = ""
+        for new_text in self.stream_generate_answer(
+                max_new_tokens=max_length,
+                temperature=temperature,
+                context_len=context_len,
+        ):
+            if new_text != stop_str:
+                response += new_text
+                yield response
+
+    def predict(
+            self,
+            query: str,
+            max_length: int = 512,
+            context_len: int = 2048,
+            temperature: float = 0.7,
+    ):
+        """Query from corpus."""
+        reference_results = []
+        if not self.enable_history:
+            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)
+            prompt = PROMPT_TEMPLATE.format(context_str=context_str, query_str=query)
+            logger.debug(f"prompt: {prompt}")
+        else:
+            prompt = query
+        self.history.append([prompt, ''])
+        response = ""
+        for new_text in self.stream_generate_answer(
+                max_new_tokens=max_length,
+                temperature=temperature,
+                context_len=context_len,
+        ):
+            response += new_text
+        response = response.strip()
+        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.")
+            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}")
+
+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("--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="maidalun1020/bce-reranker-base_v1")
+    parser.add_argument("--corpus_files", type=str, default="docs/corpus.txt")
+    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)
+    args = parser.parse_args()
+    print(args)
+    sim_model = BertSimilarity(model_name_or_path=args.sim_model_name, device=args.device)
+    m = ChatPDF(
+        similarity_model=sim_model,
+        generate_model_type=args.gen_model_type,
+        generate_model_name_or_path=args.gen_model_name,
+        lora_model_name_or_path=args.lora_model,
+        device=args.device,
+        int4=args.int4,
+        int8=args.int8,
+        chunk_size=args.chunk_size,
+        chunk_overlap=args.chunk_overlap,
+        corpus_files=args.corpus_files.split(','),
+        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(','))
+    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)