Source code for autorag.nodes.passagecompressor.longllmlingua

from typing import List, Optional

import torch
from llmlingua import PromptCompressor

from autorag.nodes.passagecompressor.base import passage_compressor_node




@passage_compressor_node
def longllmlingua(
	queries: List[str],
	contents: List[List[str]],
	scores,
	ids,
	model_name: str = "NousResearch/Llama-2-7b-hf",
	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 scores: The scores of retrieved passages.
	    Do not use in this function, so you can pass an empty list.
	:param ids: The ids of retrieved passages.
	    Do not use in this function, so you can pass an empty list.
	:param model_name: The model name to use for compression.
	    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.
	"""
	if instructions is None:
		instructions = "Given the context, please answer the final question"
	llm_lingua = PromptCompressor(
		model_name=model_name,
	)
	results = [
		llmlingua_pure(
			query, contents_, llm_lingua, instructions, target_token, **kwargs
		)
		for query, contents_ in zip(queries, contents)
	]

	del llm_lingua
	if torch.cuda.is_available():
		torch.cuda.empty_cache()

	return results





def llmlingua_pure(
	query: str,
	contents: List[str],
	llm_lingua: PromptCompressor,
	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.
	"""
	# split by "\n\n" (recommended by LongLLMLingua authors)
	new_context_texts = [c for context in contents for c in context.split("\n\n")]
	compressed_prompt = llm_lingua.compress_prompt(
		new_context_texts,
		question=query,
		instruction=instructions,
		rank_method="longllmlingua",
		target_token=target_token,
		**kwargs,
	)
	compressed_prompt_txt = compressed_prompt["compressed_prompt"]

	# separate out the question and instruction
	result = "\n\n".join(compressed_prompt_txt.split("\n\n")[1:-1])

	return result