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 "" 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)