AutoRAG 0.3.10__tar.gz → 0.3.11__tar.gz

{autorag-0.3.10 → autorag-0.3.11}/.gitignore

@@ -162,3 +162,6 @@ pytest.ini
 .DS_Store
 projects/tutorial_1
 !projects/tutorial_1/config.yaml
+
+# Visual Studio Code
+.vscode/

{autorag-0.3.10 → autorag-0.3.11}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.3.10
+Version: 0.3.11
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -219,6 +219,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: pydantic<2.10.0
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pandas>=2.1.0
 Requires-Dist: tqdm
@@ -241,11 +242,12 @@ Requires-Dist: mixedbread-ai
 Requires-Dist: llama-index-llms-bedrock
 Requires-Dist: scikit-learn
 Requires-Dist: emoji
-Requires-Dist: pymilvus
+Requires-Dist: pymilvus>=2.3.0
 Requires-Dist: chromadb>=0.5.0
 Requires-Dist: weaviate-client
 Requires-Dist: pinecone[grpc]
 Requires-Dist: couchbase
+Requires-Dist: qdrant-client
 Requires-Dist: quart
 Requires-Dist: pyngrok
 Requires-Dist: llama-index>=0.11.0

{autorag-0.3.10 → autorag-0.3.11}/AutoRAG.egg-info/SOURCES.txt

@@ -190,6 +190,7 @@ autorag/vectordb/chroma.py
 autorag/vectordb/couchbase.py
 autorag/vectordb/milvus.py
 autorag/vectordb/pinecone.py
+autorag/vectordb/qdrant.py
 autorag/vectordb/weaviate.py
 docs/Makefile
 docs/make.bat
@@ -289,6 +290,7 @@ docs/source/data_creation/chunk/chunk.md
 docs/source/data_creation/chunk/langchain_chunk.md
 docs/source/data_creation/chunk/llama_index_chunk.md
 docs/source/data_creation/legacy/legacy.md
+docs/source/data_creation/legacy/parse.md
 docs/source/data_creation/legacy/ragas.md
 docs/source/data_creation/legacy/tutorial.md
 docs/source/data_creation/parse/clova.md
@@ -315,6 +317,7 @@ docs/source/integration/vectordb/chroma.md
 docs/source/integration/vectordb/couchbase.md
 docs/source/integration/vectordb/milvus.md
 docs/source/integration/vectordb/pinecone.md
+docs/source/integration/vectordb/qdrant.md
 docs/source/integration/vectordb/vectordb.md
 docs/source/integration/vectordb/weaviate.md
 docs/source/nodes/index.md
@@ -374,7 +377,8 @@ projects/tutorial_1/config.yaml
 sample_config/chunk/chunk_full.yaml
 sample_config/chunk/chunk_ko.yaml
 sample_config/chunk/simple_chunk.yaml
-sample_config/parse/parse_full.yaml
+sample_config/parse/all_files_full.yaml
+sample_config/parse/file_types_full.yaml
 sample_config/parse/parse_hybird.yaml
 sample_config/parse/parse_ko.yaml
 sample_config/parse/parse_multimodal.yaml
@@ -535,6 +539,7 @@ tests/autorag/vectordb/test_chroma.py
 tests/autorag/vectordb/test_couchbase.py
 tests/autorag/vectordb/test_milvus.py
 tests/autorag/vectordb/test_pinecone.py
+tests/autorag/vectordb/test_qdrant.py
 tests/autorag/vectordb/test_weaviate.py
 tests/resources/README.md
 tests/resources/corpus_data_sample.parquet
@@ -555,9 +560,20 @@ tests/resources/data_creation/raw_dir/sample2.txt
 tests/resources/data_creation/raw_dir/sample3.txt
 tests/resources/parse_data/all_files/baseball_1.pdf
 tests/resources/parse_data/all_files/csv_sample.csv
+tests/resources/parse_data/all_files_full/baseball_1.pdf
+tests/resources/parse_data/all_files_full/csv_sample.csv
+tests/resources/parse_data/all_files_full/html_sample.html
+tests/resources/parse_data/all_files_full/json_sample.json
+tests/resources/parse_data/all_files_full/markdown_sample.md
+tests/resources/parse_data/all_files_full/xml_sample.xml
 tests/resources/parse_data/clova_data/result_sample.json
 tests/resources/parse_data/clova_data/result_table.txt
 tests/resources/parse_data/clova_data/result_text.txt
+tests/resources/parse_data/config/all_files.yaml
+tests/resources/parse_data/config/lack_full_parse.yaml
+tests/resources/parse_data/config/lack_simple_parse.yaml
+tests/resources/parse_data/config/perfect_full_parse.yaml
+tests/resources/parse_data/config/perfect_simple_parse.yaml
 tests/resources/parse_data/csv_data/csv_sample.csv
 tests/resources/parse_data/eng_text/baseball_1.pdf
 tests/resources/parse_data/eng_text/baseball_2.pdf

{autorag-0.3.10 → autorag-0.3.11}/AutoRAG.egg-info/requires.txt

@@ -1,3 +1,4 @@
+pydantic<2.10.0
 numpy<2.0.0
 pandas>=2.1.0
 tqdm
@@ -20,11 +21,12 @@ mixedbread-ai
 llama-index-llms-bedrock
 scikit-learn
 emoji
-pymilvus
+pymilvus>=2.3.0
 chromadb>=0.5.0
 weaviate-client
 pinecone[grpc]
 couchbase
+qdrant-client
 quart
 pyngrok
 llama-index>=0.11.0

{autorag-0.3.10 → autorag-0.3.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.3.10
+Version: 0.3.11
 Summary: Automatically Evaluate RAG pipelines with your own data. Find optimal structure for new RAG product.
 Author-email: Marker-Inc <vkehfdl1@gmail.com>
 License:                                  Apache License
@@ -219,6 +219,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
+Requires-Dist: pydantic<2.10.0
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pandas>=2.1.0
 Requires-Dist: tqdm
@@ -241,11 +242,12 @@ Requires-Dist: mixedbread-ai
 Requires-Dist: llama-index-llms-bedrock
 Requires-Dist: scikit-learn
 Requires-Dist: emoji
-Requires-Dist: pymilvus
+Requires-Dist: pymilvus>=2.3.0
 Requires-Dist: chromadb>=0.5.0
 Requires-Dist: weaviate-client
 Requires-Dist: pinecone[grpc]
 Requires-Dist: couchbase
+Requires-Dist: qdrant-client
 Requires-Dist: quart
 Requires-Dist: pyngrok
 Requires-Dist: llama-index>=0.11.0

autorag-0.3.11/autorag/VERSION

@@ -0,0 +1 @@
+0.3.11

{autorag-0.3.10 → autorag-0.3.11}/autorag/__init__.py

@@ -5,7 +5,7 @@ import sys
 from random import random
 from typing import List, Any
 
-from llama_index.core import MockEmbedding
+from llama_index.core.embeddings.mock_embed_model import MockEmbedding
 from llama_index.core.base.llms.types import CompletionResponse
 from llama_index.core.llms.mock import MockLLM
 from llama_index.llms.bedrock import Bedrock

{autorag-0.3.10 → autorag-0.3.11}/autorag/data/parse/base.py

@@ -3,6 +3,7 @@ import logging
 from datetime import datetime
 from glob import glob
 from typing import Tuple, List, Optional
+import os
 
 from autorag.utils import result_to_dataframe
 from autorag.data.utils.util import get_file_metadata
@@ -14,7 +15,10 @@ def parser_node(func):
 	@functools.wraps(func)
 	@result_to_dataframe(["texts", "path", "page", "last_modified_datetime"])
 	def wrapper(
-		data_path_glob: str, parse_method: Optional[str] = None, **kwargs
+		data_path_glob: str,
+		file_type: str,
+		parse_method: Optional[str] = None,
+		**kwargs,
 	) -> Tuple[List[str], List[str], List[int], List[datetime]]:
 		logger.info(f"Running parser - {func.__name__} module...")
 
@@ -22,6 +26,27 @@ def parser_node(func):
 		if not data_path_list:
 			raise FileNotFoundError(f"data does not exits in {data_path_glob}")
 
+		assert file_type in [
+			"pdf",
+			"csv",
+			"json",
+			"md",
+			"html",
+			"xml",
+			"all_files",
+		], f"search type {file_type} is not supported"
+
+		# extract only files from data_path_list based on the file_type set in the YAML file
+		data_paths = (
+			[
+				data_path
+				for data_path in data_path_list
+				if os.path.basename(data_path).split(".")[-1] == file_type
+			]
+			if file_type != "all_files"
+			else data_path_list
+		)
+
 		if func.__name__ == "langchain_parse":
 			parse_method = parse_method.lower()
 			if parse_method == "directory":
@@ -30,14 +55,14 @@ def parser_node(func):
 				folder_path = "/".join(path_split_list)
 				kwargs.update({"glob": glob_path, "path": folder_path})
 				result = func(
-					data_path_list=data_path_list, parse_method=parse_method, **kwargs
+					data_path_list=data_paths, parse_method=parse_method, **kwargs
 				)
 			else:
 				result = func(
-					data_path_list=data_path_list, parse_method=parse_method, **kwargs
+					data_path_list=data_paths, parse_method=parse_method, **kwargs
 				)
 		elif func.__name__ in ["clova_ocr", "llama_parse", "table_hybrid_parse"]:
-			result = func(data_path_list=data_path_list, **kwargs)
+			result = func(data_path_list=data_paths, **kwargs)
 		else:
 			raise ValueError(f"Unsupported module_type: {func.__name__}")
 		result = _add_last_modified_datetime(result)

autorag-0.3.11/autorag/data/parse/run.py

@@ -0,0 +1,120 @@
+import os
+from typing import List, Callable, Dict
+import pandas as pd
+from glob import glob
+
+from autorag.strategy import measure_speed
+from autorag.data.utils.util import get_param_combinations
+
+default_map = {
+	"pdf": {
+		"file_type": "pdf",
+		"module_type": "langchain_parse",
+		"parse_method": "pdfminer",
+	},
+	"csv": {
+		"file_type": "csv",
+		"module_type": "langchain_parse",
+		"parse_method": "csv",
+	},
+	"md": {
+		"file_type": "md",
+		"module_type": "langchain_parse",
+		"parse_method": "unstructuredmarkdown",
+	},
+	"html": {
+		"file_type": "html",
+		"module_type": "langchain_parse",
+		"parse_method": "bshtml",
+	},
+	"xml": {
+		"file_type": "xml",
+		"module_type": "langchain_parse",
+		"parse_method": "unstructuredxml",
+	},
+}
+
+
+def run_parser(
+	modules: List[Callable],
+	module_params: List[Dict],
+	data_path_glob: str,
+	project_dir: str,
+	all_files: bool,
+):
+	if not all_files:
+		# Set the parsing module to default if it is a file type in paths but not set in YAML.
+		data_path_list = glob(data_path_glob)
+		if not data_path_list:
+			raise FileNotFoundError(f"data does not exits in {data_path_glob}")
+
+		file_types = set(
+			[os.path.basename(data_path).split(".")[-1] for data_path in data_path_list]
+		)
+		set_file_types = set([module["file_type"] for module in module_params])
+
+		# create a list of only those file_types that are in file_types but not in set_file_types
+		missing_file_types = list(file_types - set_file_types)
+		if list(set_file_types - file_types):
+			raise ValueError(
+				f"File types {list(set_file_types - file_types)} are not in the data path."
+			)
+
+		if missing_file_types:
+			add_modules_list = []
+			for missing_file_type in missing_file_types:
+				if missing_file_type == "json":
+					raise ValueError(
+						"JSON file type must have a jq_schema so you must set it in the YAML file."
+					)
+
+				add_modules_list.append(default_map[missing_file_type])
+
+			add_modules, add_params = get_param_combinations(add_modules_list)
+			modules.extend(add_modules)
+			module_params.extend(add_params)
+
+	results, execution_times = zip(
+		*map(
+			lambda x: measure_speed(x[0], data_path_glob=data_path_glob, **x[1]),
+			zip(modules, module_params),
+		)
+	)
+	average_times = list(map(lambda x: x / len(results[0]), execution_times))
+
+	# save results to parquet files
+	if all_files:
+		filepaths = list(
+			map(
+				lambda x: os.path.join(project_dir, f"{x}.parquet"), range(len(modules))
+			)
+		)
+	else:
+		filepaths = list(
+			map(
+				lambda x: os.path.join(project_dir, f"{x['file_type']}.parquet"),
+				module_params,
+			)
+		)
+	list(map(lambda x: x[0].to_parquet(x[1], index=False), zip(results, filepaths)))
+	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,
+		}
+	)
+	summary_df.to_csv(os.path.join(project_dir, "summary.csv"), index=False)
+
+	# concat all parquet files here if not all_files.
+	if not all_files:
+		dataframes = [pd.read_parquet(file) for file in filepaths]
+		combined_df = pd.concat(dataframes, ignore_index=True)
+		combined_df.to_parquet(
+			os.path.join(project_dir, "parsed_result.parquet"), index=False
+		)
+
+	return summary_df

{autorag-0.3.10 → autorag-0.3.11}/autorag/data/qa/query/llama_gen_query.py

@@ -4,7 +4,7 @@ from typing import Dict, List
 from llama_index.core.base.llms.base import BaseLLM
 from llama_index.core.base.llms.types import ChatResponse, ChatMessage, MessageRole
 
-from autorag.data.qa.query.prompt import QUERY_GEN_PROMPT
+from autorag.data.qa.query.prompt import QUERY_GEN_PROMPT, QUERY_GEN_PROMPT_EXTRA
 
 
 async def llama_index_generate_base(
@@ -68,3 +68,15 @@ async def custom_query_gen(
 	messages: List[ChatMessage],
 ) -> Dict:
 	return await llama_index_generate_base(row, llm, messages)
+
+
+# Experimental feature: can only use factoid_single_hop
+async def multiple_queries_gen(
+	row: Dict,
+	llm: BaseLLM,
+	lang: str = "en",
+	n: int = 3,
+) -> Dict:
+	_messages = QUERY_GEN_PROMPT["factoid_single_hop"][lang]
+	_messages[0].content += QUERY_GEN_PROMPT_EXTRA["multiple_queries"][lang].format(n=n)
+	return await llama_index_generate_base(row, llm, _messages)

{autorag-0.3.10 → autorag-0.3.11}/autorag/data/qa/query/prompt.py

@@ -179,3 +179,12 @@ Two-hop question (using Document 2): ヌエヴォ·ラレドのテコロテス·
 		],
 	},
 }
+
+# Experimental feature
+QUERY_GEN_PROMPT_EXTRA = {
+	"multiple_queries": {
+		"en": "\nAdditional instructions:\n  - Please make {n} questions.",
+		"ko": "\n추가 지침:\n  - 질문은 {n}개를 만드세요.",
+		"ja": "\n追加指示:\n  - 質問を{n}個作成してください。",
+	}
+}

{autorag-0.3.10 → autorag-0.3.11}/autorag/data/qa/schema.py

@@ -1,5 +1,6 @@
 import logging
 from typing import Callable, Optional, Dict, Awaitable, Any, Tuple, List
+import uuid
 import pandas as pd
 from autorag.utils.util import process_batch, get_event_loop, fetch_contents
 
@@ -137,6 +138,11 @@ class QA:
 		loop = get_event_loop()
 		tasks = [fn(qa_dict, **kwargs) for qa_dict in qa_dicts]
 		results = loop.run_until_complete(process_batch(tasks, batch_size))
+
+		# Experimental feature
+		if fn.__name__ == "multiple_queries_gen":
+			return self._process_multiple_queries_gen(results)
+
 		return QA(pd.DataFrame(results), self.linked_corpus)
 
 	def batch_filter(
@@ -299,3 +305,18 @@ class QA:
 			path: corpus_df[corpus_df["path"] == path]
 			for path in corpus_df["path"].unique()
 		}
+
+	# Experimental feature
+	def _process_multiple_queries_gen(self, results: List[Dict]) -> "QA":
+		data = []
+		for result in results:
+			queries = result["query"].split("\n")
+			for query in queries:
+				new_result = {
+					key: (str(uuid.uuid4()) if key == "qid" else result[key])
+					for key in result.keys()
+				}
+				new_result["query"] = query
+				data.append(new_result)
+		df = pd.DataFrame(data)
+		return QA(df, self.linked_corpus)

{autorag-0.3.10 → autorag-0.3.11}/autorag/nodes/generator/vllm.py

@@ -37,17 +37,21 @@ class Vllm(BaseGenerator):
 	def __del__(self):
 		try:
 			import torch
+			import contextlib
 
 			if torch.cuda.is_available():
 				from vllm.distributed.parallel_state import (
-					destroy_model_parallel,
+					destroy_model_parallel, destroy_distributed_environment
 				)
 
 				destroy_model_parallel()
+				destroy_distributed_environment()
+				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.distributed.destroy_process_group()
 				torch.cuda.synchronize()
 		except ImportError:
 			del self.vllm_model

{autorag-0.3.10 → autorag-0.3.11}/autorag/parser.py

@@ -14,7 +14,7 @@ class Parser:
 		self.data_path_glob = data_path_glob
 		self.project_dir = project_dir if project_dir is not None else os.getcwd()
 
-	def start_parsing(self, yaml_path: str):
+	def start_parsing(self, yaml_path: str, all_files: bool = False):
 		if not os.path.exists(self.project_dir):
 			os.makedirs(self.project_dir)
 
@@ -32,5 +32,6 @@ class Parser:
 			module_params=input_params,
 			data_path_glob=self.data_path_glob,
 			project_dir=self.project_dir,
+			all_files=all_files,
 		)
 		logger.info("Parsing Done!")

{autorag-0.3.10 → autorag-0.3.11}/autorag/vectordb/__init__.py

@@ -18,6 +18,8 @@ def get_support_vectordb(vectordb_name: str):
 		"Pinecone": ("autorag.vectordb.pinecone", "Pinecone"),
 		"couchbase": ("autorag.vectordb.couchbase", "Couchbase"),
 		"Couchbase": ("autorag.vectordb.couchbase", "Couchbase"),
+		"qdrant": ("autorag.vectordb.qdrant", "Qdrant"),
+		"Qdrant": ("autorag.vectordb.qdrant", "Qdrant"),
 	}
 	return dynamically_find_function(vectordb_name, support_vectordb)
 

{autorag-0.3.10 → autorag-0.3.11}/autorag/vectordb/milvus.py

@@ -1,5 +1,5 @@
 import logging
-from typing import List, Tuple, Optional
+from typing import Any, Dict, List, Tuple, Optional
 
 from pymilvus import (
 	DataType,
@@ -24,13 +24,15 @@ class Milvus(BaseVectorStore):
 		embedding_model: str,
 		collection_name: str,
 		embedding_batch: int = 100,
-		similarity_metric: str = "cosine",
+		similarity_metric: str = "l2",
+		index_type: str = "IVF_FLAT",
 		uri: str = "http://localhost:19530",
 		db_name: str = "",
 		token: str = "",
 		user: str = "",
 		password: str = "",
 		timeout: Optional[float] = None,
+  		params: Dict[str, Any] = {},	
 	):
 		super().__init__(embedding_model, similarity_metric, embedding_batch)
 
@@ -45,6 +47,9 @@ class Milvus(BaseVectorStore):
 		)
 		self.collection_name = collection_name
 		self.timeout = timeout
+		self.params = params
+		self.index_type = index_type
+  
 		# Set Collection
 		if not utility.has_collection(collection_name, timeout=timeout):
 			# Get the dimension of the embeddings
@@ -66,6 +71,14 @@ class Milvus(BaseVectorStore):
 			schema = CollectionSchema(fields=[pk, field])
 
 			self.collection = Collection(name=self.collection_name, schema=schema)
+			index_params = {
+				"metric_type": self.similarity_metric.upper(),
+				"index_type": self.index_type.upper(),
+				"params": self.params,
+			}
+			self.collection.create_index(
+				field_name="vector", index_params=index_params, timeout=self.timeout
+			)
 		else:
 			self.collection = Collection(name=self.collection_name)
 
@@ -88,15 +101,6 @@ class Milvus(BaseVectorStore):
 
 		self.collection.flush(timeout=self.timeout)
 
-		index_params = {
-			"index_type": "IVF_FLAT",
-			"metric_type": self.similarity_metric.upper(),
-			"params": {},
-		}  # TODO : add params
-		self.collection.create_index(
-			field_name="vector", index_params=index_params, timeout=self.timeout
-		)
-
 	async def query(
 		self, queries: List[str], top_k: int, **kwargs
 	) -> Tuple[List[List[str]], List[List[float]]]: