AutoRAG 0.0.0__py3-none-any.whl

autorag/nodes/passagereranker/jina.py

@@ -0,0 +1,115 @@
+import os
+from typing import List, Tuple
+
+import aiohttp
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import get_event_loop, process_batch, result_to_dataframe
+
+JINA_API_URL = "https://api.jina.ai/v1/rerank"
+
+
+class JinaReranker(BasePassageReranker):
+	def __init__(self, project_dir: str, api_key: str = None, *args, **kwargs):
+		"""
+		Initialize Jina rerank node.
+
+		:param project_dir: The project directory path.
+		:param api_key: The API key for Jina rerank.
+		You can set it in the environment variable JINAAI_API_KEY.
+		Or, you can directly set it on the config YAML file using this parameter.
+		Default is env variable "JINAAI_API_KEY".
+		:param kwargs: Extra arguments that are not affected
+		"""
+		super().__init__(project_dir)
+		if api_key is None:
+			api_key = os.getenv("JINAAI_API_KEY", None)
+			if api_key is None:
+				raise ValueError(
+					"API key is not provided."
+					"You can set it as an argument or as an environment variable 'JINAAI_API_KEY'"
+				)
+		self.session = aiohttp.ClientSession(loop=get_event_loop())
+		self.session.headers.update(
+			{"Authorization": f"Bearer {api_key}", "Accept-Encoding": "identity"}
+		)
+
+	def __del__(self):
+		self.session.close()
+		del self.session
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, _, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.pop("top_k")
+		batch = kwargs.pop("batch", 8)
+		model = kwargs.pop("model", "jina-reranker-v1-base-en")
+		return self._pure(queries, contents, ids, top_k, model, batch)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		model: str = "jina-reranker-v1-base-en",
+		batch: int = 8,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents with Jina rerank models.
+		You can get the API key from https://jina.ai/reranker and set it in the environment variable JINAAI_API_KEY.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		:param ids_list: The list of lists of ids retrieved from the initial ranking
+		:param top_k: The number of passages to be retrieved
+		:param model: The model name for Cohere rerank.
+		    You can choose between "jina-reranker-v1-base-en" and "jina-colbert-v1-en".
+		    Default is "jina-reranker-v1-base-en".
+		:param batch: The number of queries to be processed in a batch
+		:return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		tasks = [
+			jina_reranker_pure(
+				self.session, query, contents, ids, top_k=top_k, model=model
+			)
+			for query, contents, ids in zip(queries, contents_list, ids_list)
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(process_batch(tasks, batch))
+
+		content_result, id_result, score_result = zip(*results)
+
+		return list(content_result), list(id_result), list(score_result)
+
+
+async def jina_reranker_pure(
+	session,
+	query: str,
+	contents: List[str],
+	ids: List[str],
+	top_k: int,
+	model: str = "jina-reranker-v1-base-en",
+) -> Tuple[List[str], List[str], List[float]]:
+	async with session.post(
+		JINA_API_URL,
+		json={
+			"query": query,
+			"documents": contents,
+			"model": model,
+			"top_n": top_k,
+		},
+	) as resp:
+		resp_json = await resp.json()
+		if "results" not in resp_json:
+			raise RuntimeError(f"Invalid response from Jina API: {resp_json['detail']}")
+
+		results = resp_json["results"]
+		indices = list(map(lambda x: x["index"], results))
+		score_result = list(map(lambda x: x["relevance_score"], results))
+		id_result = list(map(lambda x: ids[x], indices))
+		content_result = list(map(lambda x: contents[x], indices))
+
+		return content_result, id_result, score_result

autorag/nodes/passagereranker/koreranker.py

@@ -0,0 +1,136 @@
+from typing import List, Tuple
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import (
+	make_batch,
+	sort_by_scores,
+	flatten_apply,
+	select_top_k,
+	result_to_dataframe,
+	empty_cuda_cache,
+)
+
+
+class KoReranker(BasePassageReranker):
+	def __init__(self, project_dir: str, *args, **kwargs):
+		super().__init__(project_dir)
+		try:
+			import torch
+			from transformers import AutoModelForSequenceClassification, AutoTokenizer
+		except ImportError:
+			raise ImportError("For using KoReranker, please install torch first.")
+
+		model_path = "Dongjin-kr/ko-reranker"
+		self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+		self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
+		self.model.eval()
+		# Determine the device to run the model on (GPU if available, otherwise CPU)
+		self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+		self.model.to(self.device)
+
+	def __del__(self):
+		del self.model
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, _, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.pop("top_k")
+		batch = kwargs.pop("batch", 64)
+		return self._pure(queries, contents, ids, top_k, batch)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		batch: int = 64,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents based on their relevance to a query using ko-reranker.
+		ko-reranker is a reranker based on korean (https://huggingface.co/Dongjin-kr/ko-reranker).
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		:param ids_list: The list of lists of ids retrieved from the initial ranking
+		:param top_k: The number of passages to be retrieved
+		:param batch: The number of queries to be processed in a batch
+		    Default is 64.
+		:return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		nested_list = [
+			list(map(lambda x: [query, x], content_list))
+			for query, content_list in zip(queries, contents_list)
+		]
+		scores_nps = flatten_apply(
+			koreranker_run_model,
+			nested_list,
+			model=self.model,
+			batch_size=batch,
+			tokenizer=self.tokenizer,
+			device=self.device,
+		)
+
+		rerank_scores = list(
+			map(
+				lambda scores: exp_normalize(np.array(scores)).astype(float), scores_nps
+			)
+		)
+
+		df = pd.DataFrame(
+			{
+				"contents": contents_list,
+				"ids": ids_list,
+				"scores": rerank_scores,
+			}
+		)
+		df[["contents", "ids", "scores"]] = df.apply(
+			sort_by_scores, axis=1, result_type="expand"
+		)
+		results = select_top_k(df, ["contents", "ids", "scores"], top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)
+
+
+def koreranker_run_model(input_texts, model, tokenizer, device, batch_size: int):
+	try:
+		import torch
+	except ImportError:
+		raise ImportError("For using KoReranker, please install torch first.")
+	batch_input_texts = make_batch(input_texts, batch_size)
+	results = []
+	for batch_texts in batch_input_texts:
+		inputs = tokenizer(
+			batch_texts,
+			padding=True,
+			truncation=True,
+			return_tensors="pt",
+			max_length=512,
+		)
+		inputs = inputs.to(device)
+		with torch.no_grad():
+			scores = (
+				model(**inputs, return_dict=True)
+				.logits.view(
+					-1,
+				)
+				.float()
+			)
+			scores_np = scores.cpu().numpy()
+			results.extend(scores_np)
+	return results
+
+
+def exp_normalize(x):
+	b = x.max()
+	y = np.exp(x - b)
+	return y / y.sum()

autorag/nodes/passagereranker/mixedbreadai.py

@@ -0,0 +1,126 @@
+import os
+from typing import List, Tuple
+
+import pandas as pd
+from mixedbread_ai.client import AsyncMixedbreadAI
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import (
+	result_to_dataframe,
+	get_event_loop,
+	process_batch,
+	pop_params,
+)
+
+
+class MixedbreadAIReranker(BasePassageReranker):
+	def __init__(
+		self,
+		project_dir: str,
+		*args,
+		**kwargs,
+	):
+		"""
+		Initialize mixedbread-ai rerank node.
+
+		:param project_dir: The project directory path.
+		:param api_key: The API key for MixedbreadAI rerank.
+		    You can set it in the environment variable MXBAI_API_KEY.
+		    Or, you can directly set it on the config YAML file using this parameter.
+		    Default is env variable "MXBAI_API_KEY".
+		:param kwargs: Extra arguments that are not affected
+		"""
+		super().__init__(project_dir)
+		api_key = kwargs.pop("api_key", None)
+		api_key = os.getenv("MXBAI_API_KEY", None) if api_key is None else api_key
+		if api_key is None:
+			raise KeyError(
+				"Please set the API key for Mixedbread AI rerank in the environment variable MXBAI_API_KEY "
+				"or directly set it on the config YAML file."
+			)
+		self.client = AsyncMixedbreadAI(api_key=api_key)
+
+	def __del__(self):
+		del self.client
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, scores, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.pop("top_k")
+		batch = kwargs.pop("batch", 8)
+		model = kwargs.pop("model", "mixedbread-ai/mxbai-rerank-large-v1")
+		rerank_params = pop_params(self.client.reranking, kwargs)
+		return self._pure(queries, contents, ids, top_k, model, batch, **rerank_params)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		model: str = "mixedbread-ai/mxbai-rerank-large-v1",
+		batch: int = 8,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents with mixedbread-ai rerank models.
+		You can get the API key from https://www.mixedbread.ai/api-reference#quick-start-guide and set it in the environment variable MXBAI_API_KEY.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		:param ids_list: The list of lists of ids retrieved from the initial ranking
+		:param top_k: The number of passages to be retrieved
+		:param model: The model name for mixedbread-ai rerank.
+			You can choose between "mixedbread-ai/mxbai-rerank-large-v1", "mixedbread-ai/mxbai-rerank-base-v1" and "mixedbread-ai/mxbai-rerank-xsmall-v1".
+			Default is "mixedbread-ai/mxbai-rerank-large-v1".
+		:param batch: The number of queries to be processed in a batch
+		        :return: Tuple of lists containing the reranked contents, ids, and scores
+		"""
+		tasks = [
+			mixedbreadai_rerank_pure(
+				self.client, query, contents, ids, top_k=top_k, model=model
+			)
+			for query, contents, ids in zip(queries, contents_list, ids_list)
+		]
+		loop = get_event_loop()
+		results = loop.run_until_complete(process_batch(tasks, batch))
+
+		content_result, id_result, score_result = zip(*results)
+
+		return list(content_result), list(id_result), list(score_result)
+
+
+async def mixedbreadai_rerank_pure(
+	client: AsyncMixedbreadAI,
+	query: str,
+	documents: List[str],
+	ids: List[str],
+	top_k: int,
+	model: str = "mixedbread-ai/mxbai-rerank-large-v1",
+) -> Tuple[List[str], List[str], List[float]]:
+	"""
+	Rerank a list of contents with mixedbread-ai rerank models.
+
+	:param client: The mixedbread-ai client to use for reranking
+	:param query: The query to use for reranking
+	:param documents: The list of contents to rerank
+	:param ids: The list of ids corresponding to the documents
+	:param top_k: The number of passages to be retrieved
+	:param model: The model name for mixedbread-ai rerank.
+	    You can choose between "mixedbread-ai/mxbai-rerank-large-v1" and "mixedbread-ai/mxbai-rerank-base-v1".
+	    Default is "mixedbread-ai/mxbai-rerank-large-v1".
+	:return: Tuple of lists containing the reranked contents, ids, and scores
+	"""
+
+	results = await client.reranking(
+		query=query,
+		input=documents,
+		top_k=top_k,
+		model=model,
+	)
+	reranked_scores: List[float] = list(map(lambda x: x.score, results.data))
+	reranked_scores_float = list(map(float, reranked_scores))
+	indices = list(map(lambda x: x.index, results.data))
+	reranked_contents = list(map(lambda x: documents[x], indices))
+	reranked_ids: List[str] = list(map(lambda i: ids[i], indices))
+	return reranked_contents, reranked_ids, reranked_scores_float

autorag/nodes/passagereranker/monot5.py

@@ -0,0 +1,190 @@
+from itertools import chain
+from typing import List, Tuple
+
+import pandas as pd
+
+from autorag.nodes.passagereranker.base import BasePassageReranker
+from autorag.utils.util import (
+	make_batch,
+	sort_by_scores,
+	flatten_apply,
+	select_top_k,
+	result_to_dataframe,
+	pop_params,
+	empty_cuda_cache,
+)
+
+prediction_tokens = {
+	"castorini/monot5-base-msmarco": ["▁false", "▁true"],
+	"castorini/monot5-base-msmarco-10k": ["▁false", "▁true"],
+	"castorini/monot5-large-msmarco": ["▁false", "▁true"],
+	"castorini/monot5-large-msmarco-10k": ["▁false", "▁true"],
+	"castorini/monot5-base-med-msmarco": ["▁false", "▁true"],
+	"castorini/monot5-3b-med-msmarco": ["▁false", "▁true"],
+	"castorini/monot5-3b-msmarco-10k": ["▁false", "▁true"],
+	"unicamp-dl/mt5-base-en-msmarco": ["▁no", "▁yes"],
+	"unicamp-dl/ptt5-base-pt-msmarco-10k-v2": ["▁não", "▁sim"],
+	"unicamp-dl/ptt5-base-pt-msmarco-100k-v2": ["▁não", "▁sim"],
+	"unicamp-dl/ptt5-base-en-pt-msmarco-100k-v2": ["▁não", "▁sim"],
+	"unicamp-dl/mt5-base-en-pt-msmarco-v2": ["▁no", "▁yes"],
+	"unicamp-dl/mt5-base-mmarco-v2": ["▁no", "▁yes"],
+	"unicamp-dl/mt5-base-en-pt-msmarco-v1": ["▁no", "▁yes"],
+	"unicamp-dl/mt5-base-mmarco-v1": ["▁no", "▁yes"],
+	"unicamp-dl/ptt5-base-pt-msmarco-10k-v1": ["▁não", "▁sim"],
+	"unicamp-dl/ptt5-base-pt-msmarco-100k-v1": ["▁não", "▁sim"],
+	"unicamp-dl/ptt5-base-en-pt-msmarco-10k-v1": ["▁não", "▁sim"],
+	"unicamp-dl/mt5-3B-mmarco-en-pt": ["▁", "▁true"],
+	"unicamp-dl/mt5-13b-mmarco-100k": ["▁", "▁true"],
+}
+
+
+class MonoT5(BasePassageReranker):
+	def __init__(
+		self,
+		project_dir: str,
+		model_name: str = "castorini/monot5-3b-msmarco-10k",
+		*args,
+		**kwargs,
+	):
+		"""
+		Initialize the MonoT5 reranker.
+
+		:param project_dir: The project directory
+		:param model_name: The name of the MonoT5 model to use for reranking
+			Note: default model name is 'castorini/monot5-3b-msmarco-10k'
+				If there is a '/' in the model name parameter,
+				when we create the file to store the results, the path will be twisted because of the '/'.
+				Therefore, it will be received as '_' instead of '/'.
+		:param kwargs: The extra arguments for the MonoT5 reranker
+		"""
+		super().__init__(project_dir)
+		try:
+			import torch
+			from transformers import T5Tokenizer, T5ForConditionalGeneration
+		except ImportError:
+			raise ImportError("For using MonoT5 Reranker, please install torch first.")
+		# replace '_' to '/'
+		if "_" in model_name:
+			model_name = model_name.replace("_", "/")
+		# Load the tokenizer and model from the pre-trained MonoT5 model
+		self.tokenizer = T5Tokenizer.from_pretrained(model_name)
+		model_params = pop_params(T5ForConditionalGeneration.from_pretrained, kwargs)
+		self.model = T5ForConditionalGeneration.from_pretrained(
+			model_name, **model_params
+		).eval()
+
+		# Determine the device to run the model on (GPU if available, otherwise CPU)
+		self.device = "cuda" if torch.cuda.is_available() else "cpu"
+		self.model.to(self.device)
+
+		token_false, token_true = prediction_tokens[model_name]
+		self.token_false_id = self.tokenizer.convert_tokens_to_ids(token_false)
+		self.token_true_id = self.tokenizer.convert_tokens_to_ids(token_true)
+
+	def __del__(self):
+		del self.model
+		del self.tokenizer
+		empty_cuda_cache()
+		super().__del__()
+
+	@result_to_dataframe(["retrieved_contents", "retrieved_ids", "retrieve_scores"])
+	def pure(self, previous_result: pd.DataFrame, *args, **kwargs):
+		queries, contents, _, ids = self.cast_to_run(previous_result)
+		top_k = kwargs.get("top_k", 3)
+		batch = kwargs.get("batch", 64)
+		return self._pure(queries, contents, ids, top_k, batch)
+
+	def _pure(
+		self,
+		queries: List[str],
+		contents_list: List[List[str]],
+		ids_list: List[List[str]],
+		top_k: int,
+		batch: int = 64,
+	) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
+		"""
+		Rerank a list of contents based on their relevance to a query using MonoT5.
+
+		:param queries: The list of queries to use for reranking
+		:param contents_list: The list of lists of contents to rerank
+		:param ids_list: The list of lists of ids retrieved from the initial ranking
+		:param top_k: The number of passages to be retrieved
+
+		:param batch: The number of queries to be processed in a batch
+		:return: tuple of lists containing the reranked contents, ids, and scores
+		"""
+		# Retrieve the tokens used by the model to represent false and true predictions
+
+		nested_list = [
+			list(map(lambda x: [f"Query: {query} Document: {x}"], content_list))
+			for query, content_list in zip(queries, contents_list)
+		]
+
+		rerank_scores = flatten_apply(
+			monot5_run_model,
+			nested_list,
+			model=self.model,
+			batch_size=batch,
+			tokenizer=self.tokenizer,
+			device=self.device,
+			token_false_id=self.token_false_id,
+			token_true_id=self.token_true_id,
+		)
+
+		df = pd.DataFrame(
+			{
+				"contents": contents_list,
+				"ids": ids_list,
+				"scores": rerank_scores,
+			}
+		)
+		df[["contents", "ids", "scores"]] = df.apply(
+			sort_by_scores, axis=1, result_type="expand"
+		)
+		results = select_top_k(df, ["contents", "ids", "scores"], top_k)
+
+		return (
+			results["contents"].tolist(),
+			results["ids"].tolist(),
+			results["scores"].tolist(),
+		)
+
+
+def monot5_run_model(
+	input_texts,
+	model,
+	batch_size: int,
+	tokenizer,
+	device,
+	token_false_id,
+	token_true_id,
+):
+	try:
+		import torch
+	except ImportError:
+		raise ImportError("For using MonoT5 Reranker, please install torch first.")
+	batch_input_texts = make_batch(input_texts, batch_size)
+	results = []
+	for batch_texts in batch_input_texts:
+		flattened_batch_texts = list(chain.from_iterable(batch_texts))
+		input_encodings = tokenizer(
+			flattened_batch_texts,
+			padding=True,
+			truncation=True,
+			max_length=512,
+			return_tensors="pt",
+		).to(device)
+		with torch.no_grad():
+			outputs = model.generate(
+				input_ids=input_encodings["input_ids"],
+				attention_mask=input_encodings["attention_mask"],
+				output_scores=True,
+				return_dict_in_generate=True,
+			)
+
+		# Extract logits for the 'false' and 'true' tokens from the model's output
+		logits = outputs.scores[-1][:, [token_false_id, token_true_id]]
+		# Calculate the softmax probability of the 'true' token
+		probs = torch.nn.functional.softmax(logits, dim=-1)[:, 1]
+		results.extend(probs.tolist())
+	return results