AutoRAG 0.0.0__py3-none-any.whl

autorag/deploy/swagger.yml

@@ -0,0 +1,202 @@
+openapi: 3.0.0
+info:
+  title: Example API
+  version: 1.0.1
+paths:
+  /v1/run:
+    post:
+      summary: Run a query and get generated text with retrieved passages
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                query:
+                  type: string
+                  description: The query string
+                result_column:
+                  type: string
+                  description: The result column name
+                  default: generated_texts
+              required:
+                - query
+      responses:
+        '200':
+          description: Successful response
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  type:
+                    type: string
+                    enum:
+                      - generated_text
+                      - retrieved_passage
+                    description: |
+                      When the type is "generated_text", only "generated_text" is returned.
+                      The other fields are None. When the type is "retrieved_passage", only
+                      "retrieved_passage" and "passage_index" are returned. The other fields are None.
+                  generated_text:
+                    type: string
+                    nullable: true
+                    description: |
+                      The generated text, only present when "type" is "generated_text".
+                  retrieved_passage:
+                    type: object
+                    nullable: true
+                    properties:
+                      content:
+                        type: string
+                      doc_id:
+                        type: string
+                      filepath:
+                        type: string
+                        nullable: true
+                      file_page:
+                        type: integer
+                        nullable: true
+                      start_idx:
+                        type: integer
+                        nullable: true
+                      end_idx:
+                        type: integer
+                        nullable: true
+                  passage_index:
+                    type: integer
+                    nullable: true
+                    description: |
+                      The index of the retrieved passage, only present when "type" is "retrieved_passage".
+                required:
+                  - type
+                oneOf:
+                  - required:
+                      - generated_text
+                  - required:
+                      - retrieved_passage
+                      - passage_index
+  /v1/retrieve:
+    post:
+      summary: Retrieve documents based on a query
+      operationId: runRetrieveOnly
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                query:
+                  type: string
+                  description: The query string to retrieve documents.
+              required:
+                - query
+      responses:
+        '200':
+          description: Successful retrieval of documents
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  passages:
+                    type: array
+                    items:
+                      type: object
+                      properties:
+                        doc_id:
+                          type: string
+                          description: The unique identifier for the document.
+                        content:
+                          type: string
+                          description: The content of the retrieved document.
+                        score:
+                          type: number
+                          format: float
+                          description: The score of the retrieved document.
+        '400':
+          description: Invalid request due to missing query parameter
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  error:
+                    type: string
+                    description: Error message explaining the issue.
+  /v1/stream:
+    post:
+      summary: Stream generated text with retrieved passages
+      description: >
+        This endpoint streams the generated text line by line. The `retrieved_passage`
+        is sent first, followed by the `result` streamed incrementally.
+      requestBody:
+        required: true
+        content:
+          application/json:
+            schema:
+              type: object
+              properties:
+                query:
+                  type: string
+                  description: The query string
+                result_column:
+                  type: string
+                  description: The result column name
+                  default: generated_texts
+              required:
+                - query
+      responses:
+        '200':
+          description: Successful response with streaming
+          content:
+            text/event-stream:
+              schema:
+                type: object
+                properties:
+                  result:
+                    oneOf:
+                      - type: string
+                      - type: array
+                        items:
+                          type: string
+                    description: The result text or list of texts (streamed line by line)
+                  retrieved_passage:
+                    type: array
+                    items:
+                      type: object
+                      properties:
+                        content:
+                          type: string
+                        doc_id:
+                          type: string
+                        filepath:
+                          type: string
+                        nullable: true
+                        file_page:
+                          type: integer
+                          nullable: true
+                        start_idx:
+                          type: integer
+                          nullable: true
+                        end_idx:
+                          type: integer
+                          nullable: true
+
+  /version:
+    get:
+      summary: Get the API version
+      description: Returns the current version of the API as a string.
+      responses:
+        '200':
+          description: Successful response
+          content:
+            application/json:
+              schema:
+                type: object
+                properties:
+                  version:
+                    type: string
+                    description: The version of the API

autorag/embedding/base.py

@@ -0,0 +1,144 @@
+import logging
+import sys
+
+from random import random
+from typing import List, Union, Dict
+
+from llama_index.core.embeddings.mock_embed_model import MockEmbedding
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index.embeddings.openai import OpenAIEmbeddingModelType
+from llama_index.embeddings.ollama import OllamaEmbedding
+from langchain_openai.embeddings import OpenAIEmbeddings
+from llama_index.embeddings.openai_like import OpenAILikeEmbedding
+
+from autorag import LazyInit
+from autorag.embedding.vllm import VllmEmbedding
+
+logger = logging.getLogger("AutoRAG")
+
+
+class MockEmbeddingRandom(MockEmbedding):
+	"""Mock embedding with random vectors."""
+
+	def _get_vector(self) -> List[float]:
+		return [random() for _ in range(self.embed_dim)]
+
+
+embedding_models = {
+	# llama index
+	"openai": LazyInit(
+		OpenAIEmbedding
+	),  # default model is OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002
+	"openai_embed_3_large": LazyInit(
+		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_LARGE
+	),
+	"openai_embed_3_small": LazyInit(
+		OpenAIEmbedding, model_name=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL
+	),
+	"mock": LazyInit(MockEmbeddingRandom, embed_dim=768),
+	# langchain
+	"openai_langchain": LazyInit(OpenAIEmbeddings),
+	"ollama": LazyInit(OllamaEmbedding),
+	# openai like
+	"openai_like": LazyInit(OpenAILikeEmbedding),
+	"vllm": LazyInit(VllmEmbedding),
+}
+
+try:
+	# you can use your own model in this way.
+	from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+	embedding_models["huggingface_baai_bge_small"] = LazyInit(
+		HuggingFaceEmbedding, model_name="BAAI/bge-small-en-v1.5"
+	)
+	embedding_models["huggingface_cointegrated_rubert_tiny2"] = LazyInit(
+		HuggingFaceEmbedding, model_name="cointegrated/rubert-tiny2"
+	)
+	embedding_models["huggingface_all_mpnet_base_v2"] = LazyInit(
+		HuggingFaceEmbedding,
+		model_name="sentence-transformers/all-mpnet-base-v2",
+		max_length=512,
+	)
+	embedding_models["huggingface_bge_m3"] = LazyInit(
+		HuggingFaceEmbedding, model_name="BAAI/bge-m3"
+	)
+	embedding_models["huggingface_multilingual_e5_large"] = LazyInit(
+		HuggingFaceEmbedding, model_name="intfloat/multilingual-e5-large-instruct"
+	)
+except ImportError:
+	logger.info(
+		"You are using API version of AutoRAG."
+		"To use local version, run pip install 'AutoRAG[gpu]'"
+	)
+
+
+class EmbeddingModel:
+	@staticmethod
+	def load(config: Union[str, Dict, List[Dict]]):
+		if isinstance(config, str):
+			return EmbeddingModel.load_from_str(config)
+		elif isinstance(config, dict):
+			return EmbeddingModel.load_from_dict(config)
+		elif isinstance(config, list):
+			return EmbeddingModel.load_from_list(config)
+		else:
+			raise ValueError("Invalid type of config")
+
+	@staticmethod
+	def load_from_str(name: str):
+		try:
+			return embedding_models[name]
+		except KeyError:
+			raise ValueError(f"Embedding model '{name}' is not supported")
+
+	@staticmethod
+	def load_from_list(option: List[dict]):
+		if len(option) != 1:
+			raise ValueError("Only one embedding model is supported")
+		return EmbeddingModel.load_from_dict(option[0])
+
+	@staticmethod
+	def load_from_dict(option: dict):
+		def _check_keys(target: dict):
+			if "type" not in target or "model_name" not in target:
+				raise ValueError("Both 'type' and 'model_name' must be provided")
+			if target["type"] not in [
+				"openai",
+				"huggingface",
+				"mock",
+				"ollama",
+				"vllm",
+			]:
+				raise ValueError(
+					f"Embedding model type '{target['type']}' is not supported"
+				)
+
+		def _get_huggingface_class():
+			module = sys.modules.get("llama_index.embeddings.huggingface")
+			if not module:
+				logger.info(
+					"You are using API version of AutoRAG. "
+					"To use local version, run `pip install 'AutoRAG[gpu]'`."
+				)
+				return None
+			return getattr(module, "HuggingFaceEmbedding", None)
+
+		_check_keys(option)
+
+		model_options = option
+		model_type = model_options.pop("type")
+
+		embedding_map = {
+			"openai": OpenAIEmbedding,
+			"mock": MockEmbeddingRandom,
+			"huggingface": _get_huggingface_class(),
+			"ollama": OllamaEmbedding,
+			"openai_like": OpenAILikeEmbedding,
+			"vllm": VllmEmbedding,
+		}
+
+		embedding_class = embedding_map.get(model_type)
+		if not embedding_class:
+			raise ValueError(f"Embedding model type '{model_type}' is not supported")
+
+		return LazyInit(embedding_class, **model_options)

autorag/embedding/vllm.py

@@ -0,0 +1,256 @@
+from io import BytesIO
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+from llama_index.core.bridge.pydantic import Field, PrivateAttr
+from llama_index.core.callbacks import CallbackManager
+from llama_index.core.embeddings.multi_modal_base import MultiModalEmbedding
+from llama_index.core.schema import ImageType
+from PIL import Image
+from tenacity import retry, stop_after_attempt, wait_exponential
+import atexit
+
+SUPPORT_EMBED_TYPES = ["image", "text"]
+logger = logging.getLogger("AutoRAG")
+
+
+class VllmEmbedding(MultiModalEmbedding):
+	"""
+	Vllm LLM.
+
+	This class runs a vLLM embedding model locally.
+	"""
+
+	tensor_parallel_size: Optional[int] = Field(
+		default=1,
+		description="The number of GPUs to use for distributed execution with tensor parallelism.",
+	)
+
+	trust_remote_code: Optional[bool] = Field(
+		default=True,
+		description="Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer.",
+	)
+
+	dtype: str = Field(
+		default="auto",
+		description="The data type for the model weights and activations.",
+	)
+
+	download_dir: Optional[str] = Field(
+		default=None,
+		description="Directory to download and load the weights. (Default to the default cache dir of huggingface)",
+	)
+
+	vllm_kwargs: Dict[str, Any] = Field(
+		default_factory=dict,
+		description="Holds any model parameters valid for `vllm.LLM` call not explicitly specified.",
+	)
+
+	retry_cnt: int = 3
+	wait_multiplier: int = 1
+	wait_min: int = 4
+	wait_max: int = 10
+
+	_client: Any = PrivateAttr()
+
+	_image_token_id: Union[int, None] = PrivateAttr()
+
+	def __init__(
+		self,
+		model_name: str = "facebook/opt-125m",
+		embed_batch_size: int = 128,
+		tensor_parallel_size: int = 1,
+		trust_remote_code: bool = False,
+		dtype: str = "auto",
+		download_dir: Optional[str] = None,
+		vllm_kwargs: Dict[str, Any] = {},
+		callback_manager: Optional[CallbackManager] = None,
+	) -> None:
+		callback_manager = callback_manager or CallbackManager([])
+		super().__init__(
+			model_name=model_name,
+			embed_batch_size=embed_batch_size,
+			callback_manager=callback_manager,
+		)
+		try:
+			from vllm import LLM as VLLModel
+		except ImportError:
+			raise ImportError(
+				"Could not import vllm python package. "
+				"Please install it with `pip install vllm`."
+			)
+		self._client = VLLModel(
+			model=model_name,
+			task="embed",
+			max_num_seqs=embed_batch_size,
+			tensor_parallel_size=tensor_parallel_size,
+			trust_remote_code=trust_remote_code,
+			dtype=dtype,
+			download_dir=download_dir,
+			**vllm_kwargs,
+		)
+		try:
+			self._image_token_id = (
+				self._client.llm_engine.model_config.hf_config.image_token_id
+			)
+		except AttributeError:
+			self._image_token_id = None
+
+	@classmethod
+	def class_name(cls) -> str:
+		return "VllmEmbedding"
+
+	@atexit.register
+	def close():
+		import torch
+		import gc
+
+		if torch.cuda.is_available():
+			gc.collect()
+			torch.cuda.empty_cache()
+			torch.cuda.synchronize()
+
+	@retry(
+		stop=stop_after_attempt(retry_cnt),
+		wait=wait_exponential(multiplier=wait_multiplier, min=wait_min, max=wait_max),
+		reraise=True,
+	)
+	def _embed_with_retry(
+		self, inputs: List[Union[str, BytesIO]], embed_type: str = "text"
+	) -> List[List[float]]:
+		"""
+		Generates embeddings with retry mechanism.
+
+		Args:
+		    inputs: List of texts or images to embed
+
+		Returns:
+		    List of embedding vectors
+
+		Raises:
+		    Exception: If embedding fails after retries
+
+		"""
+		try:
+			if embed_type == "image":
+				inputs = [
+					{
+						"prompt_token_ids": [self._image_token_id],
+						"multi_modal_data": {"image": x},
+					}
+					for x in inputs
+				]
+			emb = self._client.embed(inputs)
+			return [x.outputs.embedding for x in emb]
+		except Exception as e:
+			logger.warning(f"Embedding attempt failed: {e!s}")
+			raise
+
+	def _embed(
+		self, inputs: List[Union[str, BytesIO]], embed_type: str = "text"
+	) -> List[List[float]]:
+		"""
+		Generates Embeddings with input validation and retry mechanism.
+
+		Args:
+		    sentences: Texts or Sentences to embed
+		    prompt_name: The name of the prompt to use for encoding
+
+		Returns:
+		    List of embedding vectors
+
+		Raises:
+		    ValueError: If any input text is invalid
+		    Exception: If embedding fails after retries
+
+		"""
+		if embed_type not in SUPPORT_EMBED_TYPES:
+			raise (ValueError("Not Implemented"))
+		return self._embed_with_retry(inputs, embed_type)
+
+	def _get_query_embedding(self, query: str) -> List[float]:
+		"""
+		Generates Embeddings for Query.
+
+		Args:
+		    query (str): Query text/sentence
+
+		Returns:
+		    List[float]: numpy array of embeddings
+
+		"""
+		return self._embed([query])[0]
+
+	async def _aget_query_embedding(self, query: str) -> List[float]:
+		"""
+		Generates Embeddings for Query Asynchronously.
+
+		Args:
+		    query (str): Query text/sentence
+
+		Returns:
+		    List[float]: numpy array of embeddings
+
+		"""
+		return self._get_query_embedding(query)
+
+	async def _aget_text_embedding(self, text: str) -> List[float]:
+		"""
+		Generates Embeddings for text Asynchronously.
+
+		Args:
+		    text (str): Text/Sentence
+
+		Returns:
+		    List[float]: numpy array of embeddings
+
+		"""
+		return self._get_text_embedding(text)
+
+	def _get_text_embedding(self, text: str) -> List[float]:
+		"""
+		Generates Embeddings for text.
+
+		Args:
+		    text (str): Text/sentences
+
+		Returns:
+		    List[float]: numpy array of embeddings
+
+		"""
+		return self._embed([text])[0]
+
+	def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
+		"""
+		Generates Embeddings for text.
+
+		Args:
+		    texts (List[str]): Texts / Sentences
+
+		Returns:
+		    List[List[float]]: numpy array of embeddings
+
+		"""
+		return self._embed(texts)
+
+	def _get_image_embedding(self, img_file_path: ImageType) -> List[float]:
+		"""Generate embedding for an image."""
+		image = Image.open(img_file_path)
+		return self._embed([image], "image")[0]
+
+	async def _aget_image_embedding(self, img_file_path: ImageType) -> List[float]:
+		"""Generate embedding for an image asynchronously."""
+		return self._get_image_embedding(img_file_path)
+
+	def _get_image_embeddings(
+		self, img_file_paths: List[ImageType]
+	) -> List[List[float]]:
+		images = [Image.open(x) for x in img_file_paths]
+		"""Generate embeddings for multiple images."""
+		return self._embed(images, "image")
+
+	async def _aget_image_embeddings(
+		self, img_file_paths: List[ImageType]
+	) -> List[List[float]]:
+		"""Generate embeddings for multiple images asynchronously."""
+		return self._get_image_embeddings(img_file_paths)

autorag/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from .retrieval import evaluate_retrieval
+from .generation import evaluate_generation
+from .retrieval_contents import evaluate_retrieval_contents

autorag/evaluation/generation.py

@@ -0,0 +1,88 @@
+import functools
+import warnings
+from typing import List, Callable, Union, Dict
+
+import pandas as pd
+
+from autorag.evaluation.metric.generation import (
+	bleu,
+	meteor,
+	rouge,
+	sem_score,
+	g_eval,
+	bert_score,
+	deepeval_faithfulness,
+)
+from autorag.evaluation.util import cast_metrics
+from autorag.schema.metricinput import MetricInput
+
+GENERATION_METRIC_FUNC_DICT = {
+	func.__name__: func
+	for func in [
+		bleu,
+		meteor,
+		rouge,
+		sem_score,
+		g_eval,
+		bert_score,
+		deepeval_faithfulness,
+	]
+}
+
+
+def evaluate_generation(
+	metric_inputs: List[MetricInput], metrics: Union[List[str], List[Dict]]
+):
+	def decorator_evaluate_generation(func: Callable):
+		@functools.wraps(func)
+		def wrapper(*args, **kwargs) -> pd.DataFrame:
+			generation_result = func(*args, **kwargs)
+			if type(generation_result) is tuple:
+				assert (
+					type(generation_result[0]) is list
+					and type(generation_result[0][0]) is str
+				), (
+					"Input func must return string list as generated answer at the first return value."
+				)
+				generated_str = generation_result[0]
+			elif type(generation_result) is list:
+				assert type(generation_result[0]) is str, (
+					"Input func must return string list as generated answer at the first return value."
+				)
+				generated_str = generation_result
+			else:
+				raise ValueError(
+					"Input func must return string list as generated answer at the first return value."
+				)
+			for metric_input, generated_text in zip(metric_inputs, generated_str):
+				metric_input.generated_texts = generated_text
+
+			metric_scores = {}
+			metric_names, metric_params = cast_metrics(metrics)
+
+			for metric_name, metric_param in zip(metric_names, metric_params):
+				if metric_name not in GENERATION_METRIC_FUNC_DICT:
+					warnings.warn(
+						f"metric {metric_name} is not in supported metrics: {GENERATION_METRIC_FUNC_DICT.keys()}"
+						f"{metric_name} will be ignored."
+					)
+				else:
+					metric_scores[metric_name] = GENERATION_METRIC_FUNC_DICT[
+						metric_name
+					](
+						metric_inputs=metric_inputs,
+						**metric_param,
+					)
+
+			metric_result_df = pd.DataFrame(metric_scores)
+			execution_result_df = pd.DataFrame({"generated_texts": generated_str})
+			if type(generation_result) is tuple:
+				execution_result_df["generated_tokens"] = generation_result[1]
+				execution_result_df["generated_log_probs"] = generation_result[2]
+
+			result_df = pd.concat([execution_result_df, metric_result_df], axis=1)
+			return result_df
+
+		return wrapper
+
+	return decorator_evaluate_generation