Source code for autorag.nodes.passagecompressor.longllmlingua
fromtypingimportList,Optionalimportpandasaspdfromautorag.nodes.passagecompressor.baseimportBasePassageCompressorfromautorag.utils.utilimportpop_params,result_to_dataframe,empty_cuda_cache# TODO: Parallel Processing Refactoring at #460
[docs]classLongLLMLingua(BasePassageCompressor):def__init__(self,project_dir:str,model_name:str="NousResearch/Llama-2-7b-hf",**kwargs):try:fromllmlinguaimportPromptCompressorexceptImportError:raiseImportError("LongLLMLingua is not installed. Please install it by running `pip install llmlingua`.")super().__init__(project_dir)model_init_params=pop_params(PromptCompressor.__init__,kwargs)self.llm_lingua=PromptCompressor(model_name=model_name,**model_init_params)def__del__(self):delself.llm_linguaempty_cuda_cache()super().__del__()
def_pure(self,queries:List[str],contents:List[List[str]],instructions:Optional[str]=None,target_token:int=300,**kwargs,)->List[str]:""" Compresses the retrieved texts using LongLLMLingua. For more information, visit https://github.com/microsoft/LLMLingua. :param queries: The queries for retrieved passages. :param contents: The contents of retrieved passages. :param model_name: The model name to use for compression. The default is "NousResearch/Llama-2-7b-hf". :param instructions: The instructions for compression. Default is None. When it is None, it will use default instructions. :param target_token: The target token for compression. Default is 300. :param kwargs: Additional keyword arguments. :return: The list of compressed texts. """ifinstructionsisNone:instructions="Given the context, please answer the final question"results=[llmlingua_pure(query,contents_,self.llm_lingua,instructions,target_token,**kwargs)forquery,contents_inzip(queries,contents)]returnresults
[docs]defllmlingua_pure(query:str,contents:List[str],llm_lingua,instructions:str,target_token:int=300,**kwargs,)->str:""" Return the compressed text. :param query: The query for retrieved passages. :param contents: The contents of retrieved passages. :param llm_lingua: The llm instance, that will be used to compress. :param instructions: The instructions for compression. :param target_token: The target token for compression. Default is 300. :param kwargs: Additional keyword arguments. :return: The compressed text. """try:fromllmlinguaimportPromptCompressorexceptImportError:raiseImportError("LongLLMLingua is not installed. Please install it by running `pip install llmlingua`.")# split by "\n\n" (recommended by LongLLMLingua authors)new_context_texts=[cforcontextincontentsforcincontext.split("\n\n")]compress_prompt_params=pop_params(PromptCompressor.compress_prompt,kwargs)compressed_prompt=llm_lingua.compress_prompt(new_context_texts,question=query,instruction=instructions,rank_method="longllmlingua",target_token=target_token,**compress_prompt_params,)compressed_prompt_txt=compressed_prompt["compressed_prompt"]# separate out the question and instructionresult="\n\n".join(compressed_prompt_txt.split("\n\n")[1:-1])returnresult