AutoRAG 0.2.6__tar.gz → 0.2.8__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.6
+Version: 0.2.8
 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
@@ -504,11 +505,11 @@ autorag run_web --trial_path your/path/to/trial_path
 
 <img width="1491" alt="web_interface" src="https://github.com/Marker-Inc-Korea/AutoRAG/assets/96727832/f6b00353-f6bb-4d8f-8740-1c264c0acbb8">
 
-# Star History
+# ⭐ Star History ⭐
 
 [![Star History Chart](https://api.star-history.com/svg?repos=Marker-Inc-Korea/AutoRAG&type=Date)](https://star-history.com/#Marker-Inc-Korea/AutoRAG&Date)
 
-# Contributors ⭐
+# ✨ Contributors ✨
 
 Thanks go to these wonderful people:
 

{autorag-0.2.6 → autorag-0.2.8}/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.8}/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.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.6
+Version: 0.2.8
 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
@@ -504,11 +505,11 @@ autorag run_web --trial_path your/path/to/trial_path
 
 <img width="1491" alt="web_interface" src="https://github.com/Marker-Inc-Korea/AutoRAG/assets/96727832/f6b00353-f6bb-4d8f-8740-1c264c0acbb8">
 
-# Star History
+# ⭐ Star History ⭐
 
 [![Star History Chart](https://api.star-history.com/svg?repos=Marker-Inc-Korea/AutoRAG&type=Date)](https://star-history.com/#Marker-Inc-Korea/AutoRAG&Date)
 
-# Contributors ⭐
+# ✨ Contributors ✨
 
 Thanks go to these wonderful people:
 

{autorag-0.2.6 → autorag-0.2.8}/README.md

@@ -234,11 +234,11 @@ autorag run_web --trial_path your/path/to/trial_path
 
 <img width="1491" alt="web_interface" src="https://github.com/Marker-Inc-Korea/AutoRAG/assets/96727832/f6b00353-f6bb-4d8f-8740-1c264c0acbb8">
 
-# Star History
+# ⭐ Star History ⭐
 
 [![Star History Chart](https://api.star-history.com/svg?repos=Marker-Inc-Korea/AutoRAG&type=Date)](https://star-history.com/#Marker-Inc-Korea/AutoRAG&Date)
 
-# Contributors ⭐
+# ✨ Contributors ✨
 
 Thanks go to these wonderful people:
 

autorag-0.2.8/autorag/VERSION

@@ -0,0 +1 @@
+0.2.8

{autorag-0.2.6 → autorag-0.2.8}/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.8/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.8/autorag/evaluation}/metric/generation.py

@@ -15,12 +15,13 @@ 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.utils.util import process_batch, openai_truncate_by_token
+from autorag.evaluation.metric.util import calculate_cosine_similarity
+from autorag.utils.util import process_batch, openai_truncate_by_token, convert_inputs_to_list
 
 
 def generation_metric(func):
     @functools.wraps(func)
+    @convert_inputs_to_list
     def wrapper(generation_gt: List[List[str]], generations: List[str], **kwargs) -> List[float]:
         """
         Compute generation metric.
@@ -39,6 +40,7 @@ def generation_metric(func):
     return wrapper
 
 
+@convert_inputs_to_list
 def huggingface_evaluate(instance, key: str,
                          generation_gt: List[List[str]], generations: List[str],
                          **kwargs) -> List[float]:
@@ -70,6 +72,7 @@ def bleu(gt: List[str], pred: str, **kwargs) -> float:
     return sacrebleu.sentence_bleu(pred, gt, **kwargs).score
 
 
+@convert_inputs_to_list
 def meteor(generation_gt: List[List[str]], generations: List[str],
            alpha: float = 0.9,
            beta: float = 3.0,
@@ -97,6 +100,7 @@ def meteor(generation_gt: List[List[str]], generations: List[str],
     return result
 
 
+@convert_inputs_to_list
 def rouge(generation_gt: List[List[str]], generations: List[str],
           rouge_type: Optional[str] = 'rougeL',
           use_stemmer: bool = False,
@@ -141,6 +145,7 @@ def rouge(generation_gt: List[List[str]], generations: List[str],
     return result
 
 
+@convert_inputs_to_list
 def sem_score(generation_gt: List[List[str]], generations: List[str],
               embedding_model: Optional[BaseEmbedding] = None,
               batch: int = 128) -> List[float]:
@@ -194,6 +199,7 @@ def sem_score(generation_gt: List[List[str]], generations: List[str],
     return result
 
 
+@convert_inputs_to_list
 def g_eval(generation_gt: List[List[str]], generations: List[str],
            metrics: Optional[List[str]] = None,
            model: str = 'gpt-4-0125-preview',
@@ -283,6 +289,7 @@ async def async_g_eval(generation_gt: List[str], pred: str,
     return sum(g_eval_scores) / len(g_eval_scores)
 
 
+@convert_inputs_to_list
 def bert_score(generation_gt: List[List[str]], generations: List[str],
                lang: str = 'en',
                batch: int = 128,

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

@@ -1,11 +1,15 @@
 import functools
 import itertools
-import math
 from typing import List
 
+import math
+
+from autorag.utils.util import convert_inputs_to_list
+
 
 def retrieval_metric(func):
     @functools.wraps(func)
+    @convert_inputs_to_list
     def wrapper(retrieval_gt: List[List[List[str]]], pred_ids: List[List[str]]) -> List[float]:
         results = []
         for gt, pred in zip(retrieval_gt, pred_ids):

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

@@ -9,11 +9,12 @@ from typing import List
 
 import numpy as np
 
-from autorag.utils.util import normalize_string
+from autorag.utils.util import normalize_string, convert_inputs_to_list
 
 
 def retrieval_contents_metric(func):
     @functools.wraps(func)
+    @convert_inputs_to_list
     def wrapper(gt_contents: List[List[str]], pred_contents: List[List[str]]) -> List[float]:
         results = []
         for gt, pred in zip(gt_contents, pred_contents):

{autorag-0.2.6/autorag/evaluate → autorag-0.2.8/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.8/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/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.8}/autorag/nodes/passagereranker/flag_embedding.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import List, Tuple, Iterable
 
 import pandas as pd
 import torch
@@ -58,7 +58,7 @@ def flag_embedding_run_model(input_texts, model, batch_size: int):
     for batch_texts in tqdm(batch_input_texts):
         with torch.no_grad():
             pred_scores = model.compute_score(sentence_pairs=batch_texts)
-        if batch_size == 1:
+        if batch_size == 1 or not isinstance(pred_scores, Iterable):
             results.append(pred_scores)
         else:
             results.extend(pred_scores)

{autorag-0.2.6 → autorag-0.2.8}/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]