AutoRAG 0.2.6__tar.gz → 0.2.7__tar.gz

{autorag-0.2.6 → autorag-0.2.7}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.6
+Version: 0.2.7
 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
@@ -210,16 +210,17 @@ Keywords: RAG,AutoRAG,autorag,rag-evaluation,evaluation,rag-auto,AutoML,AutoML-R
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Information Technology
 Classifier: Intended Audience :: Science/Research
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy
+Requires-Dist: numpy<2.0.0
 Requires-Dist: pandas
 Requires-Dist: tqdm
 Requires-Dist: tiktoken>=0.7.0
@@ -248,11 +249,11 @@ Requires-Dist: bert_score
 Requires-Dist: sentence-transformers
 Requires-Dist: FlagEmbedding
 Requires-Dist: ragas
-Requires-Dist: ray
 Requires-Dist: kiwipiepy
 Requires-Dist: llmlingua
 Requires-Dist: llama-index>=0.10.1
 Requires-Dist: llama-index-core>=0.10.1
+Requires-Dist: llama-index-readers-file
 Requires-Dist: llama-index-embeddings-openai
 Requires-Dist: llama-index-embeddings-huggingface
 Requires-Dist: llama-index-llms-openai

{autorag-0.2.6 → autorag-0.2.7}/AutoRAG.egg-info/SOURCES.txt

@@ -38,20 +38,20 @@ autorag/data/qacreation/ragas.py
 autorag/data/qacreation/simple.py
 autorag/data/utils/__init__.py
 autorag/data/utils/util.py
-autorag/evaluate/__init__.py
-autorag/evaluate/generation.py
-autorag/evaluate/retrieval.py
-autorag/evaluate/retrieval_contents.py
-autorag/evaluate/util.py
-autorag/evaluate/metric/__init__.py
-autorag/evaluate/metric/generation.py
-autorag/evaluate/metric/retrieval.py
-autorag/evaluate/metric/retrieval_contents.py
-autorag/evaluate/metric/util.py
-autorag/evaluate/metric/g_eval_prompts/coh_detailed.txt
-autorag/evaluate/metric/g_eval_prompts/con_detailed.txt
-autorag/evaluate/metric/g_eval_prompts/flu_detailed.txt
-autorag/evaluate/metric/g_eval_prompts/rel_detailed.txt
+autorag/evaluation/__init__.py
+autorag/evaluation/generation.py
+autorag/evaluation/retrieval.py
+autorag/evaluation/retrieval_contents.py
+autorag/evaluation/util.py
+autorag/evaluation/metric/__init__.py
+autorag/evaluation/metric/generation.py
+autorag/evaluation/metric/retrieval.py
+autorag/evaluation/metric/retrieval_contents.py
+autorag/evaluation/metric/util.py
+autorag/evaluation/metric/g_eval_prompts/coh_detailed.txt
+autorag/evaluation/metric/g_eval_prompts/con_detailed.txt
+autorag/evaluation/metric/g_eval_prompts/flu_detailed.txt
+autorag/evaluation/metric/g_eval_prompts/rel_detailed.txt
 autorag/nodes/__init__.py
 autorag/nodes/generator/__init__.py
 autorag/nodes/generator/base.py
@@ -170,8 +170,8 @@ docs/source/api_spec/autorag.data.corpus.rst
 docs/source/api_spec/autorag.data.qacreation.rst
 docs/source/api_spec/autorag.data.rst
 docs/source/api_spec/autorag.data.utils.rst
-docs/source/api_spec/autorag.evaluate.metric.rst
-docs/source/api_spec/autorag.evaluate.rst
+docs/source/api_spec/autorag.evaluation.metric.rst
+docs/source/api_spec/autorag.evaluation.rst
 docs/source/api_spec/autorag.nodes.generator.rst
 docs/source/api_spec/autorag.nodes.passageaugmenter.rst
 docs/source/api_spec/autorag.nodes.passagecompressor.rst

{autorag-0.2.6 → autorag-0.2.7}/AutoRAG.egg-info/requires.txt

@@ -1,4 +1,4 @@
-numpy
+numpy<2.0.0
 pandas
 tqdm
 tiktoken>=0.7.0
@@ -27,11 +27,11 @@ bert_score
 sentence-transformers
 FlagEmbedding
 ragas
-ray
 kiwipiepy
 llmlingua
 llama-index>=0.10.1
 llama-index-core>=0.10.1
+llama-index-readers-file
 llama-index-embeddings-openai
 llama-index-embeddings-huggingface
 llama-index-llms-openai

{autorag-0.2.6 → autorag-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.6
+Version: 0.2.7
 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
@@ -210,16 +210,17 @@ Keywords: RAG,AutoRAG,autorag,rag-evaluation,evaluation,rag-auto,AutoML,AutoML-R
 Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Information Technology
 Classifier: Intended Audience :: Science/Research
-Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy
+Requires-Dist: numpy<2.0.0
 Requires-Dist: pandas
 Requires-Dist: tqdm
 Requires-Dist: tiktoken>=0.7.0
@@ -248,11 +249,11 @@ Requires-Dist: bert_score
 Requires-Dist: sentence-transformers
 Requires-Dist: FlagEmbedding
 Requires-Dist: ragas
-Requires-Dist: ray
 Requires-Dist: kiwipiepy
 Requires-Dist: llmlingua
 Requires-Dist: llama-index>=0.10.1
 Requires-Dist: llama-index-core>=0.10.1
+Requires-Dist: llama-index-readers-file
 Requires-Dist: llama-index-embeddings-openai
 Requires-Dist: llama-index-embeddings-huggingface
 Requires-Dist: llama-index-llms-openai

autorag-0.2.7/autorag/VERSION

@@ -0,0 +1 @@
+0.2.7

{autorag-0.2.6 → autorag-0.2.7}/autorag/dashboard.py

@@ -77,7 +77,7 @@ def node_view(node_dir: str):
         'float': NumberFormatter(format='0.000'),
         'bool': BooleanFormatter(),
     }
-    first_df = pd.read_parquet(os.path.join(node_dir, '0.parquet'))
+    first_df = pd.read_parquet(os.path.join(node_dir, '0.parquet'), engine='pyarrow')
 
     each_module_df_widget = pn.widgets.Tabulator(pd.DataFrame(columns=first_df.columns), name='Module DataFrame',
                                                  formatters=bokeh_formatters,
@@ -87,7 +87,7 @@ def node_view(node_dir: str):
         if event.column == 'detail':
             filename = summary_df['filename'].iloc[event.row]
             filepath = os.path.join(node_dir, filename)
-            each_module_df = pd.read_parquet(filepath)
+            each_module_df = pd.read_parquet(filepath, engine='pyarrow')
             each_module_df_widget.value = each_module_df
 
     df_widget = pn.widgets.Tabulator(summary_df, name='Summary DataFrame', formatters=bokeh_formatters,

{autorag-0.2.6/autorag/evaluate → autorag-0.2.7/autorag/evaluation}/generation.py

@@ -4,8 +4,8 @@ from typing import List, Callable, Union, Dict
 
 import pandas as pd
 
-from autorag.evaluate.metric.generation import bleu, meteor, rouge, sem_score, g_eval, bert_score
-from autorag.evaluate.util import cast_metrics
+from autorag.evaluation.metric.generation import bleu, meteor, rouge, sem_score, g_eval, bert_score
+from autorag.evaluation.util import cast_metrics
 
 GENERATION_METRIC_FUNC_DICT = {func.__name__: func for func in
                                [bleu, meteor, rouge, sem_score, g_eval, bert_score]}

{autorag-0.2.6/autorag/evaluate → autorag-0.2.7/autorag/evaluation}/metric/generation.py

@@ -15,7 +15,7 @@ from rouge_score import tokenizers
 from rouge_score.rouge_scorer import RougeScorer
 
 from autorag import embedding_models
-from autorag.evaluate.metric.util import calculate_cosine_similarity
+from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.utils.util import process_batch, openai_truncate_by_token
 
 

{autorag-0.2.6/autorag/evaluate → autorag-0.2.7/autorag/evaluation}/retrieval.py

@@ -4,9 +4,10 @@ from typing import List, Callable, Any, Tuple, Optional, Union, Dict
 
 import pandas as pd
 
-from autorag.evaluate.metric import (retrieval_recall, retrieval_precision, retrieval_f1, retrieval_ndcg, retrieval_mrr,
-                                     retrieval_map)
-from autorag.evaluate.util import cast_metrics
+from autorag.evaluation.metric import (retrieval_recall, retrieval_precision, retrieval_f1, retrieval_ndcg,
+                                       retrieval_mrr,
+                                       retrieval_map)
+from autorag.evaluation.util import cast_metrics
 
 RETRIEVAL_METRIC_FUNC_DICT = {func.__name__: func for func in
                               [retrieval_recall, retrieval_precision, retrieval_f1, retrieval_ndcg, retrieval_mrr,

{autorag-0.2.6/autorag/evaluate → autorag-0.2.7/autorag/evaluation}/retrieval_contents.py

@@ -3,7 +3,7 @@ from typing import List, Callable, Any, Tuple
 
 import pandas as pd
 
-from autorag.evaluate.metric import retrieval_token_f1, retrieval_token_precision, retrieval_token_recall
+from autorag.evaluation.metric import retrieval_token_f1, retrieval_token_precision, retrieval_token_recall
 
 
 def evaluate_retrieval_contents(retrieval_gt: List[List[str]], metrics: List[str]):

{autorag-0.2.6 → autorag-0.2.7}/autorag/evaluator.py

@@ -56,8 +56,8 @@ class Evaluator:
             raise ValueError(f"QA data path {qa_data_path} is not a parquet file.")
         if not corpus_data_path.endswith('.parquet'):
             raise ValueError(f"Corpus data path {corpus_data_path} is not a parquet file.")
-        self.qa_data = pd.read_parquet(qa_data_path)
-        self.corpus_data = pd.read_parquet(corpus_data_path)
+        self.qa_data = pd.read_parquet(qa_data_path, engine='pyarrow')
+        self.corpus_data = pd.read_parquet(corpus_data_path, engine='pyarrow')
         self.qa_data = cast_qa_dataset(self.qa_data)
         self.corpus_data = cast_corpus_dataset(self.corpus_data)
         self.project_dir = project_dir if project_dir is not None else os.getcwd()
@@ -184,7 +184,7 @@ class Evaluator:
     def _load_node_lines(yaml_path: str) -> Dict[str, List[Node]]:
         if not os.path.exists(yaml_path):
             raise ValueError(f"YAML file {yaml_path} does not exist.")
-        with open(yaml_path, 'r') as stream:
+        with open(yaml_path, 'r', encoding='utf-8') as stream:
             try:
                 yaml_dict = yaml.safe_load(stream)
             except yaml.YAMLError as exc:
@@ -299,7 +299,7 @@ class Evaluator:
 
             previous_node_dir = os.path.join(trial_path, previous_node_line, previous_node)
             best_file_pattern = f'{previous_node_dir}/best_*.parquet'
-            previous_result = pd.read_parquet(glob.glob(best_file_pattern)[0])
+            previous_result = pd.read_parquet(glob.glob(best_file_pattern)[0], engine='pyarrow')
         return previous_result
 
     @staticmethod

{autorag-0.2.6 → autorag-0.2.7}/autorag/node_line.py

@@ -38,7 +38,7 @@ def run_node_line(nodes: List[Node],
         qa_path = os.path.join(project_dir, "data", "qa.parquet")
         if not os.path.exists(qa_path):
             raise ValueError(f"qa.parquet does not exist in {qa_path}.")
-        previous_result = pd.read_parquet(qa_path)
+        previous_result = pd.read_parquet(qa_path, engine='pyarrow')
 
     summary_lst = []
     for node in nodes:

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/generator/run.py

@@ -4,8 +4,8 @@ from typing import Callable, List, Dict, Union
 
 import pandas as pd
 
-from autorag.evaluate import evaluate_generation
-from autorag.evaluate.util import cast_metrics
+from autorag.evaluation import evaluate_generation
+from autorag.evaluation.util import cast_metrics
 from autorag.strategy import measure_speed, filter_by_threshold, select_best
 
 
@@ -35,7 +35,7 @@ def run_generator_node(modules: List[Callable],
     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"))
+    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.")
     generation_gt = list(map(lambda x: x.tolist(), qa_data['generation_gt'].tolist()))
@@ -51,7 +51,8 @@ def run_generator_node(modules: List[Callable],
     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, generation_gt, strategies.get('metrics')), results))
+    results = list(
+        map(lambda result: evaluate_generator_node(result, generation_gt, strategies.get('metrics')), results))
 
     # save results to folder
     filepaths = list(map(lambda x: os.path.join(node_dir, f'{x}.parquet'), range(len(modules))))

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passageaugmenter/base.py

@@ -1,5 +1,4 @@
 import functools
-import itertools
 import logging
 import os
 from pathlib import Path
@@ -7,13 +6,12 @@ from typing import List, Union, Tuple
 
 import numpy as np
 import pandas as pd
-import torch
 
 from autorag import embedding_models
-from autorag.evaluate.metric.util import calculate_cosine_similarity
+from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.utils import (result_to_dataframe, validate_qa_dataset, fetch_contents, sort_by_scores,
                            validate_corpus_dataset, cast_corpus_dataset)
-from autorag.utils.util import reconstruct_list, filter_dict_keys, select_top_k
+from autorag.utils.util import filter_dict_keys, select_top_k, embedding_query_content
 
 logger = logging.getLogger("AutoRAG")
 
@@ -37,7 +35,7 @@ def passage_augmenter_node(func):
         assert "retrieved_ids" in previous_result.columns, "previous_result must have retrieved_ids column."
         ids = previous_result["retrieved_ids"].tolist()
 
-        corpus_df = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"))
+        corpus_df = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"), engine='pyarrow')
         validate_corpus_dataset(corpus_df)
 
         # get top_k
@@ -59,8 +57,9 @@ def passage_augmenter_node(func):
 
             # set embedding model for getting scores
             embedding_model_str = kwargs.pop("embedding_model", 'openai')
+            embedding_model = embedding_models[embedding_model_str]
             query_embeddings, contents_embeddings = embedding_query_content(queries, augmented_contents,
-                                                                            embedding_model_str, batch=128)
+                                                                            embedding_model, batch=128)
 
             # get scores from calculated cosine similarity
             augmented_scores = [
@@ -87,22 +86,3 @@ def passage_augmenter_node(func):
         return results['contents'].tolist(), results['ids'].tolist(), results['scores'].tolist()
 
     return wrapper
-
-
-def embedding_query_content(queries: List[str], contents_list: List[List[str]],
-                            embedding_model: str, batch: int = 128):
-    embedding_model = embedding_models[embedding_model]
-
-    # Embedding using batch
-    embedding_model.embed_batch_size = batch
-    query_embeddings = embedding_model.get_text_embedding_batch(queries)
-
-    content_lengths = list(map(len, contents_list))
-    content_embeddings_flatten = embedding_model.get_text_embedding_batch(list(
-        itertools.chain.from_iterable(contents_list)))
-    content_embeddings = reconstruct_list(content_embeddings_flatten, content_lengths)
-
-    del embedding_model
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    return query_embeddings, content_embeddings

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passageaugmenter/run.py

@@ -20,7 +20,7 @@ def run_passage_augmenter_node(modules: List[Callable],
     if not os.path.exists(node_line_dir):
         os.makedirs(node_line_dir)
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
-    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))
+    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"), engine='pyarrow')
     retrieval_gt = qa_df['retrieval_gt'].tolist()
     retrieval_gt = [[[str(uuid) for uuid in sub_array] if sub_array.size > 0 else [] for sub_array in inner_array]
                     for inner_array in retrieval_gt]

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagecompressor/run.py

@@ -5,7 +5,7 @@ from typing import Callable, List, Dict
 
 import pandas as pd
 
-from autorag.evaluate.metric import retrieval_token_recall, retrieval_token_precision, retrieval_token_f1
+from autorag.evaluation.metric import retrieval_token_recall, retrieval_token_precision, retrieval_token_f1
 from autorag.strategy import measure_speed, filter_by_threshold, select_best
 from autorag.utils.util import fetch_contents
 
@@ -40,8 +40,8 @@ def run_passage_compressor_node(modules: List[Callable],
         os.makedirs(save_dir)
 
     # make retrieval contents gt
-    qa_data = pd.read_parquet(os.path.join(data_dir, "qa.parquet"))
-    corpus_data = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"))
+    qa_data = pd.read_parquet(os.path.join(data_dir, "qa.parquet"), engine='pyarrow')
+    corpus_data = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"), engine='pyarrow')
     # check qa_data have retrieval_gt
     assert all(len(x[0]) > 0 for x in qa_data['retrieval_gt'].tolist()), \
         "Can't use passage compressor if you don't have retrieval gt values in QA dataset."

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagefilter/base.py

@@ -39,7 +39,7 @@ def passage_filter_node(func):
         ids = previous_result["retrieved_ids"].tolist()
 
         if func.__name__ == 'recency_filter':
-            corpus_df = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"))
+            corpus_df = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"), engine='pyarrow')
             metadatas = fetch_contents(corpus_df, ids, column_name='metadata')
             times = [[time['last_modified_datetime'] for time in time_list] for time_list in metadatas]
             filtered_contents, filtered_ids, filtered_scores \

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagefilter/run.py

@@ -30,7 +30,7 @@ def run_passage_filter_node(modules: List[Callable],
     if not os.path.exists(node_line_dir):
         os.makedirs(node_line_dir)
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
-    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))
+    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"), engine='pyarrow')
     retrieval_gt = qa_df['retrieval_gt'].tolist()
     retrieval_gt = [[[str(uuid) for uuid in sub_array] if sub_array.size > 0 else [] for sub_array in inner_array]
                     for inner_array in retrieval_gt]

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagefilter/similarity_percentile_cutoff.py

@@ -3,7 +3,8 @@ from typing import List, Tuple, Optional
 import numpy as np
 import torch.cuda
 
-from autorag.evaluate.metric.util import calculate_cosine_similarity
+from autorag import embedding_models
+from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import passage_filter_node
 from autorag.nodes.passagefilter.similarity_threshold_cutoff import embedding_query_content
 
@@ -31,7 +32,13 @@ def similarity_percentile_cutoff(queries: List[str], contents_list: List[List[st
         Default is 128.
     :return: Tuple of lists containing the filtered contents, ids, and scores
     """
-    query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model, batch)
+    if embedding_model is None:
+        embedding_model = embedding_models['openai']
+    else:
+        embedding_model = embedding_models[embedding_model]
+
+    query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model,
+                                                                   batch)
 
     results = list(map(lambda x: similarity_percentile_cutoff_pure(x[0], x[1], x[2], x[3], x[4], percentile),
                        zip(query_embeddings, content_embeddings, contents_list, ids_list, scores_list)))

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagefilter/similarity_threshold_cutoff.py

@@ -1,13 +1,12 @@
-import itertools
 from typing import List, Tuple, Optional
 
 import numpy as np
 import torch.cuda
 
 from autorag import embedding_models
-from autorag.evaluate.metric.util import calculate_cosine_similarity
+from autorag.evaluation.metric.util import calculate_cosine_similarity
 from autorag.nodes.passagefilter.base import passage_filter_node
-from autorag.utils.util import reconstruct_list
+from autorag.utils.util import embedding_query_content
 
 
 @passage_filter_node
@@ -33,6 +32,11 @@ def similarity_threshold_cutoff(queries: List[str], contents_list: List[List[str
         Default is 128.
     :return: Tuple of lists containing the filtered contents, ids, and scores
     """
+    if embedding_model is None:
+        embedding_model = embedding_models['openai']
+    else:
+        embedding_model = embedding_models[embedding_model]
+
     query_embeddings, content_embeddings = embedding_query_content(queries, contents_list, embedding_model, batch)
 
     remain_indices = list(map(lambda x: similarity_threshold_cutoff_pure(x[0], x[1], threshold),
@@ -68,21 +72,3 @@ def similarity_threshold_cutoff_pure(query_embedding: str,
     if len(result) > 0:
         return result
     return [np.argmax(similarities)]
-
-
-def embedding_query_content(queries: List[str], contents_list: List[List[str]],
-                            embedding_model: Optional[str] = None, batch: int = 128):
-    if embedding_model is None:
-        embedding_model = embedding_models['openai']
-    else:
-        embedding_model = embedding_models[embedding_model]
-
-    # Embedding using batch
-    embedding_model.embed_batch_size = batch
-    query_embeddings = embedding_model.get_text_embedding_batch(queries)
-
-    content_lengths = list(map(len, contents_list))
-    content_embeddings_flatten = embedding_model.get_text_embedding_batch(list(
-        itertools.chain.from_iterable(contents_list)))
-    content_embeddings = reconstruct_list(content_embeddings_flatten, content_lengths)
-    return query_embeddings, content_embeddings

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagereranker/base.py

@@ -39,7 +39,7 @@ def passage_reranker_node(func):
 
         # time rerankers
         if func.__name__ == 'time_reranker':
-            corpus_df = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"))
+            corpus_df = pd.read_parquet(os.path.join(project_dir, "data", "corpus.parquet"), engine='pyarrow')
             metadatas = fetch_contents(corpus_df, ids, column_name='metadata')
             times = [[time['last_modified_datetime'] for time in time_list] for time_list in metadatas]
             reranked_contents, reranked_ids, reranked_scores \

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagereranker/run.py

@@ -34,7 +34,7 @@ def run_passage_reranker_node(modules: List[Callable],
     if not os.path.exists(node_line_dir):
         os.makedirs(node_line_dir)
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
-    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))
+    qa_df = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"), engine='pyarrow')
     retrieval_gt = qa_df['retrieval_gt'].tolist()
     retrieval_gt = [[[str(uuid) for uuid in sub_array] if sub_array.size > 0 else [] for sub_array in inner_array]
                     for inner_array in retrieval_gt]

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/passagereranker/upr.py

@@ -1,11 +1,9 @@
 import logging
-import os
-from typing import List, Tuple, Dict
+from typing import List, Tuple
 
-import numpy as np
 import pandas as pd
-import ray
 import torch
+from tqdm import tqdm
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 from autorag.nodes.passagereranker.base import passage_reranker_node
@@ -19,8 +17,7 @@ def upr(queries: List[str], contents_list: List[List[str]],
         scores_list: List[List[float]], ids_list: List[List[str]],
         top_k: int, use_bf16: bool = False,
         prefix_prompt: str = "Passage: ",
-        suffix_prompt: str = "Please write a question based on this passage.",
-        num_gpus: int = 1) \
+        suffix_prompt: str = "Please write a question based on this passage.") \
         -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
     """
     Rerank a list of contents based on their relevance to a query using UPR.
@@ -42,33 +39,20 @@ def upr(queries: List[str], contents_list: List[List[str]],
         Default is "Please write a question based on this passage.".
         The suffix prompt provides a cue or a closing instruction to the language model,
             signaling how to conclude the generated text or what format to follow at the end.
-    :param num_gpus: The number of GPUs that you want to use. Default is 1.
-        Set to 1 if you use a CPU or a single GPU.
     :return: tuple of lists containing the reranked contents, ids, and scores
     """
+    tqdm.pandas()
     df = pd.DataFrame({
         'query': queries,
         'contents': contents_list,
         'ids': ids_list,
     })
-    ds = ray.data.from_pandas(df)
 
     scorer = UPRScorer(suffix_prompt=suffix_prompt, prefix_prompt=prefix_prompt, use_bf16=use_bf16)
 
-    if torch.cuda.is_available():
-        score_batch = ds.map_batches(scorer, batch_size=1,
-                                     concurrency=num_gpus, num_gpus=1)
-    else:
-        score_batch = ds.map_batches(scorer,
-                                     batch_size=1,
-                                     concurrency=min(len(df), os.cpu_count()), num_cpus=os.cpu_count())
-    scores = score_batch.to_pandas()['output'].tolist()  # converted to a flatten list of scores
-
+    df['scores'] = df.progress_apply(lambda row: scorer.compute(query=row['query'], contents=row['contents']),
+                                     axis=1)
     del scorer
-
-    explode_df = df.explode('contents')
-    explode_df['scores'] = scores
-    df['scores'] = explode_df.groupby(level=0, sort=False)['scores'].apply(list).tolist()
     df[['contents', 'ids', 'scores']] = df.apply(lambda x: sort_by_scores(x, reverse=False), axis=1,
                                                  result_type='expand')
     results = select_top_k(df, ['contents', 'ids', 'scores'], top_k)
@@ -87,9 +71,8 @@ class UPRScorer:
         self.suffix_prompt = suffix_prompt
         self.prefix_prompt = prefix_prompt
 
-    def __call__(self, batch: Dict[str, np.ndarray]):
-        query_token = self.tokenizer(batch['query'][0], max_length=128, truncation=True, return_tensors='pt')
-        contents = batch['contents'][0]
+    def compute(self, query: str, contents: List[str]) -> List[float]:
+        query_token = self.tokenizer(query, max_length=128, truncation=True, return_tensors='pt')
         prompts = list(map(lambda content: f"{self.prefix_prompt} {content} {self.suffix_prompt}", contents))
         prompt_token_outputs = self.tokenizer(prompts, padding='longest',
                                               max_length=512,
@@ -107,7 +90,7 @@ class UPRScorer:
         log_softmax = torch.nn.functional.log_softmax(logits, dim=-1)
         nll = -log_softmax.gather(2, query_input_ids.unsqueeze(2)).squeeze(2)
         avg_nll = torch.sum(nll, dim=1)
-        return {"output": avg_nll.tolist()}
+        return avg_nll.tolist()
 
     def __del__(self):
         del self.model

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/promptmaker/run.py

@@ -6,8 +6,8 @@ from typing import List, Callable, Dict, Optional, Union
 import pandas as pd
 import tokenlog
 
-from autorag.evaluate import evaluate_generation
-from autorag.evaluate.util import cast_metrics
+from autorag.evaluation import evaluate_generation
+from autorag.evaluation.util import cast_metrics
 from autorag.strategy import measure_speed, filter_by_threshold, select_best
 from autorag.support import get_support_modules
 from autorag.utils import validate_qa_dataset
@@ -101,7 +101,7 @@ def run_prompt_maker_node(modules: List[Callable],
         generator_callables, generator_params = make_generator_callable_params(extra_strategy)
 
         # get generation_gt
-        qa_data = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))
+        qa_data = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"), engine='pyarrow')
         validate_qa_dataset(qa_data)
         generation_gt = qa_data['generation_gt'].tolist()
         generation_gt = list(map(lambda x: x.tolist(), generation_gt))

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/queryexpansion/run.py

@@ -88,7 +88,7 @@ def run_query_expansion_node(modules: List[Callable],
         retrieval_callables, retrieval_params = make_retrieval_callable_params(extra_strategy)
 
         # get retrieval_gt
-        retrieval_gt = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"))['retrieval_gt'].tolist()
+        retrieval_gt = pd.read_parquet(os.path.join(project_dir, "data", "qa.parquet"), engine='pyarrow')['retrieval_gt'].tolist()
 
         # run evaluation
         evaluation_results = list(map(lambda result: evaluate_one_query_expansion_node(

{autorag-0.2.6 → autorag-0.2.7}/autorag/nodes/retrieval/base.py

@@ -94,7 +94,7 @@ def retrieval_node(func):
             raise ValueError(f"invalid func name for using retrieval_io decorator.")
 
         # fetch data from corpus_data
-        corpus_data = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"))
+        corpus_data = pd.read_parquet(os.path.join(data_dir, "corpus.parquet"), engine='pyarrow')
         contents = fetch_contents(corpus_data, ids)
 
         return contents, ids, scores