chatpdf.py

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)
        print("\nRespuesta: ", r)