AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/generator/run.py

@@ -0,0 +1,148 @@
+import os
+import pathlib
+from typing import List, Dict, Union
+
+import pandas as pd
+
+from autorag.evaluation import evaluate_generation
+from autorag.evaluation.util import cast_metrics
+from autorag.schema.metricinput import MetricInput
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
+from autorag.utils.util import to_list
+
+
+def run_generator_node(
+	modules: List,
+	module_params: List[Dict],
+	previous_result: pd.DataFrame,
+	node_line_dir: str,
+	strategies: Dict,
+) -> pd.DataFrame:
+	"""
+	Run evaluation and select the best module among generator node results.
+	And save the results and summary to generator node directory.
+
+	:param modules: Generator modules to run.
+	:param module_params: Generator module parameters.
+	    Including node parameters, which is used for every module in this node.
+	:param previous_result: Previous result dataframe.
+	    Could be prompt maker node's result.
+	:param node_line_dir: This node line's directory.
+	:param strategies: Strategies for generator node.
+	:return: The best result dataframe.
+	    It contains previous result columns and generator node's result columns.
+	"""
+	if not os.path.exists(node_line_dir):
+		os.makedirs(node_line_dir)
+	project_dir = pathlib.PurePath(node_line_dir).parent.parent
+	node_dir = os.path.join(node_line_dir, "generator")  # node name
+	if not os.path.exists(node_dir):
+		os.makedirs(node_dir)
+	qa_data = pd.read_parquet(
+		os.path.join(project_dir, "data", "qa.parquet"), engine="pyarrow"
+	)
+	if "generation_gt" not in qa_data.columns:
+		raise ValueError("You must have 'generation_gt' column in qa.parquet.")
+
+	results, execution_times = zip(
+		*map(
+			lambda x: measure_speed(
+				x[0].run_evaluator,
+				project_dir=project_dir,
+				previous_result=previous_result,
+				**x[1],
+			),
+			zip(modules, module_params),
+		)
+	)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	# get average token usage
+	token_usages = list(map(lambda x: x["generated_tokens"].apply(len).mean(), results))
+
+	# make rows to metric_inputs
+	generation_gt = to_list(qa_data["generation_gt"].tolist())
+	queries = to_list(qa_data["query"].tolist())
+
+	metric_inputs = [
+		MetricInput(generation_gt=gen_gt, query=query)
+		for gen_gt, query in zip(generation_gt, queries)
+	]
+
+	metric_names, metric_params = cast_metrics(strategies.get("metrics"))
+	if metric_names is None or len(metric_names) <= 0:
+		raise ValueError("You must at least one metrics for generator evaluation.")
+	results = list(
+		map(
+			lambda result: evaluate_generator_node(
+				result, metric_inputs, strategies.get("metrics")
+			),
+			results,
+		)
+	)
+
+	# save results to folder
+	filepaths = list(
+		map(lambda x: os.path.join(node_dir, f"{x}.parquet"), range(len(modules)))
+	)
+	list(
+		map(lambda x: x[0].to_parquet(x[1], index=False), zip(results, filepaths))
+	)  # execute save to parquet
+	filenames = list(map(lambda x: os.path.basename(x), filepaths))
+
+	summary_df = pd.DataFrame(
+		{
+			"filename": filenames,
+			"module_name": list(map(lambda module: module.__name__, modules)),
+			"module_params": module_params,
+			"execution_time": average_times,
+			"average_output_token": token_usages,
+			**{
+				metric: list(map(lambda x: x[metric].mean(), results))
+				for metric in metric_names
+			},
+		}
+	)
+
+	# filter by strategies
+	if strategies.get("speed_threshold") is not None:
+		results, filenames = filter_by_threshold(
+			results, average_times, strategies["speed_threshold"], filenames
+		)
+	if strategies.get("token_threshold") is not None:
+		results, filenames = filter_by_threshold(
+			results, token_usages, strategies["token_threshold"], filenames
+		)
+	selected_result, selected_filename = select_best(
+		results, metric_names, filenames, strategies.get("strategy", "mean")
+	)
+	best_result = pd.concat([previous_result, selected_result], axis=1)
+
+	# add 'is_best' column at summary file
+	summary_df["is_best"] = summary_df["filename"] == selected_filename
+
+	# save files
+	summary_df.to_csv(os.path.join(node_dir, "summary.csv"), index=False)
+	best_result.to_parquet(
+		os.path.join(
+			node_dir, f"best_{os.path.splitext(selected_filename)[0]}.parquet"
+		),
+		index=False,
+	)
+	return best_result
+
+
+def evaluate_generator_node(
+	result_df: pd.DataFrame,
+	metric_inputs: List[MetricInput],
+	metrics: Union[List[str], List[Dict]],
+):
+	@evaluate_generation(metric_inputs=metric_inputs, metrics=metrics)
+	def evaluate_generation_module(df: pd.DataFrame):
+		return (
+			df["generated_texts"].tolist(),
+			df["generated_tokens"].tolist(),
+			df["generated_log_probs"].tolist(),
+		)
+
+	return evaluate_generation_module(result_df)

autorag/nodes/generator/vllm.py

@@ -0,0 +1,147 @@
+import gc
+from copy import deepcopy
+from typing import List, Tuple, Union
+
+import pandas as pd
+
+from autorag.nodes.generator.base import BaseGenerator
+from autorag.utils import result_to_dataframe
+from autorag.utils.util import pop_params, to_list, is_chat_prompt
+
+
+class Vllm(BaseGenerator):
+	def __init__(self, project_dir: str, llm: str, **kwargs):
+		super().__init__(project_dir, llm, **kwargs)
+		try:
+			from vllm import SamplingParams, LLM
+		except ImportError:
+			raise ImportError(
+				"Please install vllm library. You can install it by running `pip install vllm`."
+			)
+
+		model_from_kwargs = kwargs.pop("model", None)
+		model = llm if model_from_kwargs is None else model_from_kwargs
+
+		input_kwargs = deepcopy(kwargs)
+		sampling_params_init_params = pop_params(
+			SamplingParams.from_optional, input_kwargs
+		)
+		input_kwargs.pop("thinking", None)
+		self.vllm_model = LLM(model, **input_kwargs)
+
+		# delete not sampling param keys in the kwargs
+		kwargs_keys = list(kwargs.keys())
+		for key in kwargs_keys:
+			if key not in sampling_params_init_params:
+				kwargs.pop(key)
+
+	def __del__(self):
+		try:
+			import torch
+			import contextlib
+
+			if torch.cuda.is_available():
+				from vllm.distributed.parallel_state import (
+					destroy_model_parallel,
+					destroy_distributed_environment,
+				)
+
+				destroy_model_parallel()
+				destroy_distributed_environment()
+				if hasattr(self.vllm_model.llm_engine, "model_executor"):
+					del self.vllm_model.llm_engine.model_executor
+				del self.vllm_model
+				with contextlib.suppress(AssertionError):
+					torch.distributed.destroy_process_group()
+				gc.collect()
+				torch.cuda.empty_cache()
+				torch.cuda.synchronize()
+		except ImportError:
+			del self.vllm_model
+
+		super().__del__()
+
+	@result_to_dataframe(["generated_texts", "generated_tokens", "generated_log_probs"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		prompts = self.cast_to_run(previous_result)
+		thinking = kwargs.pop("thinking", False)
+		return self._pure(prompts, thinking=thinking, **kwargs)
+
+	def _pure(
+		self,
+		prompts: Union[List[str], List[List[dict]]],
+		thinking: bool = False,
+		**kwargs,
+	) -> Tuple[List[str], List[List[int]], List[List[float]]]:
+		"""
+		Vllm module.
+		It gets the VLLM instance and returns generated texts by the input prompt.
+		You can set logprobs to get the log probs of the generated text.
+		Default logprobs is 1.
+
+		:param prompts: A list of prompts or a list of chat prompts.
+		:param thinking: A boolean that indicates whether to think when generating text.
+			Default is False.
+			Effective when set True and using chat prompts.
+			You can learn how to use chat prompt at `chat_fstring` module documentation.
+		:param kwargs: The extra parameters for generating the text.
+		:return: A tuple of three elements.
+		    The first element is a list of generated text.
+		    The second element is a list of generated text's token ids.
+		    The third element is a list of generated text's log probs.
+		"""
+		try:
+			from vllm.outputs import RequestOutput
+			from vllm import SamplingParams, LLM
+
+			try:
+				from vllm.logprobs import SampleLogprobs  # vLLM >= 0.11
+			except ImportError:  # pragma: no cover - older vLLM fallback
+				from vllm.sequence import SampleLogprobs
+		except ImportError:
+			raise ImportError(
+				"Please install vllm library. You can install it by running `pip install vllm`."
+			)
+
+		if "logprobs" not in kwargs:
+			kwargs["logprobs"] = 1
+
+		sampling_params = pop_params(SamplingParams.from_optional, kwargs)
+		generate_params = SamplingParams(**sampling_params)
+		if is_chat_prompt(prompts):
+			chat_template_kwargs = kwargs.pop("chat_template_kwargs", {})
+			chat_template_kwargs["enable_thinking"] = thinking
+			chat_kwargs = pop_params(LLM.chat, kwargs)
+			results: List[RequestOutput] = self.vllm_model.chat(
+				prompts,
+				generate_params,
+				chat_template_kwargs=chat_template_kwargs,
+				**chat_kwargs,
+			)
+		else:
+			generate_kwargs = pop_params(LLM.generate, kwargs)
+			results: List[RequestOutput] = self.vllm_model.generate(
+				prompts, generate_params, **generate_kwargs
+			)
+		generated_texts = list(map(lambda x: x.outputs[0].text, results))
+		generated_token_ids = list(map(lambda x: x.outputs[0].token_ids, results))
+		log_probs: List[SampleLogprobs] = list(
+			map(lambda x: x.outputs[0].logprobs, results)
+		)
+		generated_log_probs = list(
+			map(
+				lambda x: list(map(lambda y: y[0][y[1]].logprob, zip(x[0], x[1]))),
+				zip(log_probs, generated_token_ids),
+			)
+		)
+		return (
+			to_list(generated_texts),
+			to_list(generated_token_ids),
+			to_list(generated_log_probs),
+		)
+
+	async def astream(self, prompt: str, **kwargs):
+		raise NotImplementedError
+
+	def stream(self, prompt: str, **kwargs):
+		raise NotImplementedError

autorag/nodes/generator/vllm_api.py

@@ -0,0 +1,191 @@
+import logging
+from typing import List, Tuple, Dict, Union
+import time
+
+import pandas as pd
+import requests
+from asyncio import to_thread
+
+from autorag.nodes.generator.base import BaseGenerator
+from autorag.utils.util import get_event_loop, process_batch, result_to_dataframe
+
+logger = logging.getLogger("AutoRAG")
+
+DEFAULT_MAX_TOKENS = 4096  # Default token limit
+
+
+class VllmAPI(BaseGenerator):
+	def __init__(
+		self,
+		project_dir,
+		llm: str,
+		uri: str,
+		max_tokens: int = None,
+		batch: int = 16,
+		*args,
+		**kwargs,
+	):
+		"""
+		VLLM API Wrapper for OpenAI-compatible chat/completions format.
+
+		:param project_dir: Project directory.
+		:param llm: Model name (e.g., LLaMA model).
+		:param uri: VLLM API server URI.
+		:param max_tokens: Maximum token limit.
+		    Default is 4096.
+		:param batch: Request batch size.
+		    Default is 16.
+		"""
+		super().__init__(project_dir, llm, *args, **kwargs)
+		assert batch > 0, "Batch size must be greater than 0."
+		self.uri = uri.rstrip("/")  # Set API URI
+		self.batch = batch
+		# Use the provided max_tokens if available, otherwise use the default
+		self.max_token_size = max_tokens if max_tokens else DEFAULT_MAX_TOKENS
+		self.max_model_len = self.get_max_model_length()
+		logger.info(f"{llm} max model length: {self.max_model_len}")
+
+	@result_to_dataframe(["generated_texts", "generated_tokens", "generated_log_probs"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		prompts = self.cast_to_run(previous_result)
+		return self._pure(prompts, **kwargs)
+
+	def _pure(
+		self,
+		prompts: Union[List[str], List[List[Dict]]],
+		truncate: bool = True,
+		**kwargs,
+	) -> Tuple[List[str], List[List[int]], List[List[float]]]:
+		"""
+		Method to call the VLLM API to generate text.
+
+		:param prompts: List of input prompts.
+		:param truncate: Whether to truncate input prompts to fit within the token limit.
+		:param kwargs: Additional options (e.g., temperature, top_p).
+		:return: Generated text, token lists, and log probability lists.
+		"""
+		if kwargs.get("logprobs") is not None:
+			kwargs.pop("logprobs")
+			logger.warning(
+				"parameter logprob does not effective. It always set to True."
+			)
+		if kwargs.get("n") is not None:
+			kwargs.pop("n")
+			logger.warning("parameter n does not effective. It always set to 1.")
+
+		if truncate:
+			prompts = list(map(lambda p: self.truncate_by_token(p), prompts))
+		loop = get_event_loop()
+		tasks = [to_thread(self.get_result, prompt, **kwargs) for prompt in prompts]
+		results = loop.run_until_complete(process_batch(tasks, self.batch))
+
+		answer_result = list(map(lambda x: x[0], results))
+		token_result = list(map(lambda x: x[1], results))
+		logprob_result = list(map(lambda x: x[2], results))
+		return answer_result, token_result, logprob_result
+
+	def truncate_by_token(self, prompt: Union[str, List[Dict]]) -> str:
+		"""
+		Function to truncate prompts to fit within the maximum token limit.
+		"""
+		if not isinstance(prompt, str):
+			content_list = [msg["content"] for msg in prompt]
+			prompt = "\n".join(content_list)
+		tokens = self.encoding_for_model(prompt)["tokens"]  # Simple tokenization
+		return self.decoding_for_model(tokens[: self.max_model_len])["prompt"]
+
+	def call_vllm_api(self, prompt: Union[str, List[Dict]], **kwargs) -> dict:
+		"""
+		Calls the VLLM API to get chat/completions responses.
+
+		:param prompt: Input prompt.
+		:param kwargs: Additional API options (e.g., temperature, max_tokens).
+		:return: API response.
+		"""
+		payload = {
+			"model": self.llm,
+			"messages": parse_prompt(prompt),
+			"temperature": kwargs.get("temperature", 0.4),
+			"max_tokens": min(
+				kwargs.get("max_tokens", self.max_token_size), self.max_token_size
+			),
+			"logprobs": True,
+			"n": 1,
+		}
+		start_time = time.time()  # Record request start time
+		response = requests.post(f"{self.uri}/v1/chat/completions", json=payload)
+		end_time = time.time()  # Record request end time
+
+		response.raise_for_status()
+		elapsed_time = end_time - start_time  # Calculate elapsed time
+		logger.info(
+			f"Request chat completions to vllm server completed in {elapsed_time:.2f} seconds"
+		)
+		return response.json()
+
+	# Additional method: abstract method implementation
+	async def astream(self, prompt: Union[str, List[Dict]], **kwargs):
+		"""
+		Asynchronous streaming method not implemented.
+		"""
+		raise NotImplementedError("astream method is not implemented for VLLM API yet.")
+
+	def stream(self, prompt: Union[str, List[Dict]], **kwargs):
+		"""
+		Synchronous streaming method not implemented.
+		"""
+		raise NotImplementedError("stream method is not implemented for VLLM API yet.")
+
+	def get_result(self, prompt: Union[str, List[Dict]], **kwargs):
+		response = self.call_vllm_api(prompt, **kwargs)
+		choice = response["choices"][0]
+		answer = choice["message"]["content"]
+
+		# Handle cases where logprobs is None
+		if choice.get("logprobs") and "content" in choice["logprobs"]:
+			logprobs = list(map(lambda x: x["logprob"], choice["logprobs"]["content"]))
+			tokens = list(
+				map(
+					lambda x: self.encoding_for_model(x["token"])["tokens"],
+					choice["logprobs"]["content"],
+				)
+			)
+		else:
+			logprobs = []
+			tokens = []
+
+		return answer, tokens, logprobs
+
+	def encoding_for_model(self, answer_piece: str):
+		payload = {
+			"model": self.llm,
+			"prompt": answer_piece,
+			"add_special_tokens": True,
+		}
+		response = requests.post(f"{self.uri}/tokenize", json=payload)
+		response.raise_for_status()
+		return response.json()
+
+	def decoding_for_model(self, tokens: list[int]):
+		payload = {
+			"model": self.llm,
+			"tokens": tokens,
+		}
+		response = requests.post(f"{self.uri}/detokenize", json=payload)
+		response.raise_for_status()
+		return response.json()
+
+	def get_max_model_length(self):
+		response = requests.get(f"{self.uri}/v1/models")
+		response.raise_for_status()
+		json_data = response.json()
+		return json_data["data"][0]["max_model_len"]
+
+
+def parse_prompt(prompt: Union[str, List[Dict]]) -> List[Dict]:
+	if isinstance(prompt, str):
+		return [{"role": "user", "content": prompt}]
+	elif isinstance(prompt, list):
+		return prompt
+	else:
+		raise ValueError("prompt must be a string or a list of dicts.")

autorag/nodes/hybridretrieval/__init__.py

@@ -0,0 +1,2 @@
+from .hybrid_cc import HybridCC
+from .hybrid_rrf import HybridRRF

autorag/nodes/hybridretrieval/base.py

@@ -0,0 +1,58 @@
+import abc
+
+import pandas as pd
+
+from autorag.nodes.retrieval.base import BaseRetrieval
+from autorag.utils import result_to_dataframe
+from autorag.utils.util import pop_params, fetch_contents
+
+
+class HybridRetrieval(BaseRetrieval, metaclass=abc.ABCMeta):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		super().__init__(project_dir)
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		previous_info = self.cast_to_run(previous_result, *args, **kwargs)
+		_pure_params = pop_params(self._pure, kwargs)
+		ids, scores = self._pure(previous_info, **_pure_params)
+		contents = fetch_contents(self.corpus_df, ids)
+		return contents, ids, scores
+
+	def cast_to_run(self, previous_result: pd.DataFrame, *args, **kwargs):
+		return hybrid_cast(previous_result)
+
+	@classmethod
+	def cast_to_run_class(cls, previous_result: pd.DataFrame):
+		return hybrid_cast(previous_result)
+
+
+def hybrid_cast(
+	previous_result: pd.DataFrame,
+):
+	assert "query" in previous_result.columns, "previous_result must have query column."
+	queries = previous_result["query"].tolist()
+
+	assert "retrieved_contents_semantic" in previous_result.columns
+	assert "retrieved_contents_lexical" in previous_result.columns
+	assert "retrieve_scores_semantic" in previous_result.columns
+	assert "retrieve_scores_lexical" in previous_result.columns
+	assert "retrieved_ids_semantic" in previous_result.columns
+	assert "retrieved_ids_lexical" in previous_result.columns
+
+	contents_semantic = previous_result["retrieved_contents_semantic"].tolist()
+	contents_lexical = previous_result["retrieved_contents_lexical"].tolist()
+	scores_semantic = previous_result["retrieve_scores_semantic"].tolist()
+	scores_lexical = previous_result["retrieve_scores_lexical"].tolist()
+	ids_semantic = previous_result["retrieved_ids_semantic"].tolist()
+	ids_lexical = previous_result["retrieved_ids_lexical"].tolist()
+
+	return {
+		"queries": queries,
+		"retrieved_contents_semantic": contents_semantic,
+		"retrieved_contents_lexical": contents_lexical,
+		"retrieve_scores_semantic": scores_semantic,
+		"retrieve_scores_lexical": scores_lexical,
+		"retrieved_ids_semantic": ids_semantic,
+		"retrieved_ids_lexical": ids_lexical,
+	}