AutoRAG 0.0.3__tar.gz → 0.0.4__tar.gz

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.0.3
+Version: 0.0.4
 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
@@ -310,7 +310,7 @@ For evaluation, you need to prepare just three files.
 
 There is a template for your evaluation data for using AutoRAG.
 Check out the evaluation data rule at [here]().
-Plus, you can get example datasets for testing AutoRAG at [here]().
+Plus, you can get example datasets for testing AutoRAG at [here](./sample_dataset).
 
 ### Evaluate your data to various RAG modules
 

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/AutoRAG.egg-info/SOURCES.txt

@@ -74,6 +74,8 @@ docs/make.bat
 docs/requirements.txt
 docs/source/conf.py
 docs/source/index.rst
+docs/source/install.md
+docs/source/tutorial.md
 docs/source/api_spec/autorag.data.rst
 docs/source/api_spec/autorag.evaluate.metric.rst
 docs/source/api_spec/autorag.evaluate.rst
@@ -89,6 +91,14 @@ docs/source/api_spec/autorag.rst
 docs/source/api_spec/autorag.schema.rst
 docs/source/api_spec/autorag.utils.rst
 docs/source/api_spec/modules.rst
+docs/source/data_creation/tutorial.md
+docs/source/modules/index.md
+docs/source/optimization/optimization.md
+docs/source/resources/samsung_sundae.jpeg
+sample_dataset/README.md
+sample_dataset/eli5/load_eli5_dataset.py
+sample_dataset/msmarco/load_msmarco_dataset.py
+sample_dataset/triviaqa/load_triviaqa_dataset.py
 tests/conftest.py
 tests/autorag/test_deploy.py
 tests/autorag/test_evaluator.py
@@ -128,6 +138,7 @@ tests/autorag/utils/test_util.py
 tests/resources/corpus_data_sample.parquet
 tests/resources/full.yaml
 tests/resources/qa_data_sample.parquet
+tests/resources/qa_test_data_sample.parquet
 tests/resources/simple.yaml
 tests/resources/test_bm25_retrieval.pkl
 tests/resources/result_project/trial.json

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.0.3
+Version: 0.0.4
 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
@@ -310,7 +310,7 @@ For evaluation, you need to prepare just three files.
 
 There is a template for your evaluation data for using AutoRAG.
 Check out the evaluation data rule at [here]().
-Plus, you can get example datasets for testing AutoRAG at [here]().
+Plus, you can get example datasets for testing AutoRAG at [here](./sample_dataset).
 
 ### Evaluate your data to various RAG modules
 

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/README.md

@@ -66,7 +66,7 @@ For evaluation, you need to prepare just three files.
 
 There is a template for your evaluation data for using AutoRAG.
 Check out the evaluation data rule at [here]().
-Plus, you can get example datasets for testing AutoRAG at [here]().
+Plus, you can get example datasets for testing AutoRAG at [here](./sample_dataset).
 
 ### Evaluate your data to various RAG modules
 

AutoRAG-0.0.4/autorag/VERSION

@@ -0,0 +1 @@
+0.0.4

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/autorag/deploy.py

@@ -27,6 +27,20 @@ def extract_node_line_names(config_dict: Dict) -> List[str]:
     return [node_line['node_line_name'] for node_line in config_dict['node_lines']]
 
 
+def extract_node_strategy(config_dict: Dict) -> Dict:
+    """
+    Extract node strategies with the given config dictionary.
+    The return value is a dictionary of node type and its strategy.
+
+    :param config_dict: The yaml configuration dict for the pipeline.
+        You can load this to access trail_folder/config.yaml.
+    :return: Key is node_type and value is strategy dict.
+    """
+    return {node['node_type']: node.get('strategy', {})
+            for node_line in config_dict['node_lines']
+            for node in node_line['nodes']}
+
+
 def summary_df_to_yaml(summary_df: pd.DataFrame, config_dict: Dict) -> Dict:
     """
     Convert trial summary dataframe to config yaml file.
@@ -41,6 +55,12 @@ def summary_df_to_yaml(summary_df: pd.DataFrame, config_dict: Dict) -> Dict:
     # summary_df columns : 'node_line_name', 'node_type', 'best_module_filename',
     #                      'best_module_name', 'best_module_params', 'best_execution_time'
     node_line_names = extract_node_line_names(config_dict)
+    node_strategies = extract_node_strategy(config_dict)
+    strategy_df = pd.DataFrame({
+        'node_type': list(node_strategies.keys()),
+        'strategy': list(node_strategies.values())
+    })
+    summary_df = summary_df.merge(strategy_df, on='node_type', how='left')
     summary_df['categorical_node_line_name'] = pd.Categorical(summary_df['node_line_name'], categories=node_line_names,
                                                               ordered=True)
     summary_df = summary_df.sort_values(by='categorical_node_line_name')
@@ -52,6 +72,7 @@ def summary_df_to_yaml(summary_df: pd.DataFrame, config_dict: Dict) -> Dict:
             'nodes': [
                 {
                     'node_type': row['node_type'],
+                    'strategy': row['strategy'],
                     'modules': [{
                         'module_type': row['best_module_name'],
                         **row['best_module_params']

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/autorag/nodes/passagecompressor/tree_summarize.py

@@ -26,7 +26,7 @@ def tree_summarize(queries: List[str],
     """
     Recursively merge retrieved texts and summarizes them in a bottom-up fashion.
     This function is a wrapper for llama_index.response_synthesizers.TreeSummarize.
-    For more information, visit https://docs.llamaindex.ai/en/stable/examples/response_synthesizers/tree_summarize.html.
+    For more information, visit https://docs.llamaindex.ai/en/latest/examples/response_synthesizers/tree_summarize.html.
 
     :param queries: The queries for retrieved passages.
     :param contents: The contents of retrieved passages.

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/autorag/nodes/retrieval/hybrid_cc.py

@@ -55,8 +55,9 @@ def hybrid_cc(
 
 def cc_pure(ids: Tuple, scores: Tuple, weights: Tuple, top_k: int) -> Tuple[
     List[str], List[float]]:
-    df = pd.concat([pd.Series(dict(zip(_id, score))) for _id, score in zip(ids, scores)], axis=1, join="inner")
+    df = pd.concat([pd.Series(dict(zip(_id, score))) for _id, score in zip(ids, scores)], axis=1)
     normalized_scores = (df - df.min()) / (df.max() - df.min())
+    normalized_scores = normalized_scores.fillna(0)
     normalized_scores['weighted_sum'] = normalized_scores.mul(weights).sum(axis=1)
     normalized_scores = normalized_scores.sort_values(by='weighted_sum', ascending=False)
     return normalized_scores.index.tolist()[:top_k], normalized_scores['weighted_sum'][:top_k].tolist()

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/autorag/nodes/retrieval/run.py

@@ -68,28 +68,37 @@ def run_retrieval_node(modules: List[Callable],
 
     # run retrieval modules except hybrid
     hybrid_module_names = ['hybrid_rrf', 'hybrid_cc']
-    non_hybrid_modules, non_hybrid_module_params = zip(*filter(lambda x: x[0].__name__ not in hybrid_module_names,
-                                                               zip(modules, module_params)))
     filename_first = 0
-    non_hybrid_results, non_hybrid_times, non_hybrid_summary_df = run_and_save(non_hybrid_modules,
-                                                                               non_hybrid_module_params, filename_first)
-    filename_first += len(non_hybrid_modules)
+    if any([module.__name__ not in hybrid_module_names for module in modules]):
+        non_hybrid_modules, non_hybrid_module_params = zip(*filter(lambda x: x[0].__name__ not in hybrid_module_names,
+                                                                   zip(modules, module_params)))
+        non_hybrid_results, non_hybrid_times, non_hybrid_summary_df = run_and_save(non_hybrid_modules,
+                                                                                   non_hybrid_module_params, filename_first)
+        filename_first += len(non_hybrid_modules)
+    else:
+        non_hybrid_results, non_hybrid_times, non_hybrid_summary_df = [], [], pd.DataFrame()
 
     if any([module.__name__ in hybrid_module_names for module in modules]):
         hybrid_modules, hybrid_module_params = zip(*filter(lambda x: x[0].__name__ in hybrid_module_names,
                                                            zip(modules, module_params)))
-        target_modules = list(map(lambda x: x.pop('target_modules'), hybrid_module_params))
-        target_filenames = list(map(lambda x: select_result_for_hybrid(save_dir, x), target_modules))
-        ids_scores = list(map(lambda x: get_ids_and_scores(save_dir, x), target_filenames))
-        target_module_params = list(map(lambda x: get_module_params(save_dir, x), target_filenames))
-        hybrid_module_params = list(map(lambda x: {**x[0], **x[1]}, zip(hybrid_module_params, ids_scores)))
-        real_hybrid_times = list(map(lambda filename: get_hybrid_execution_times(save_dir, filename), target_filenames))
-        hybrid_results, hybrid_times, hybrid_summary_df = run_and_save(hybrid_modules, hybrid_module_params,
-                                                                       filename_first)
-        filename_first += len(hybrid_modules)
-        hybrid_times = real_hybrid_times.copy()
-        hybrid_summary_df['execution_time'] = hybrid_times
-        hybrid_summary_df = edit_summary_df_params(hybrid_summary_df, target_modules, target_module_params)
+        if all(['target_module_params' in x for x in hybrid_module_params]):
+            # If target_module_params are already given, run hybrid retrieval directly
+            hybrid_results, hybrid_times, hybrid_summary_df = run_and_save(hybrid_modules, hybrid_module_params,
+                                                                           filename_first)
+            filename_first += len(hybrid_modules)
+        else:
+            target_modules = list(map(lambda x: x.pop('target_modules'), hybrid_module_params))
+            target_filenames = list(map(lambda x: select_result_for_hybrid(save_dir, x), target_modules))
+            ids_scores = list(map(lambda x: get_ids_and_scores(save_dir, x), target_filenames))
+            target_module_params = list(map(lambda x: get_module_params(save_dir, x), target_filenames))
+            hybrid_module_params = list(map(lambda x: {**x[0], **x[1]}, zip(hybrid_module_params, ids_scores)))
+            real_hybrid_times = list(map(lambda filename: get_hybrid_execution_times(save_dir, filename), target_filenames))
+            hybrid_results, hybrid_times, hybrid_summary_df = run_and_save(hybrid_modules, hybrid_module_params,
+                                                                           filename_first)
+            filename_first += len(hybrid_modules)
+            hybrid_times = real_hybrid_times.copy()
+            hybrid_summary_df['execution_time'] = hybrid_times
+            hybrid_summary_df = edit_summary_df_params(hybrid_summary_df, target_modules, target_module_params)
     else:
         hybrid_results, hybrid_times, hybrid_summary_df = [], [], pd.DataFrame()
 
@@ -161,11 +170,11 @@ def select_result_for_hybrid(node_dir: str, target_modules: Tuple) -> List[str]:
     return best_filenames
 
 
-def get_module_params(node_dir: str, filenames: List[str]) -> List[Dict]:
+def get_module_params(node_dir: str, filenames: List[str]) -> Tuple[Dict]:
     summary_df = load_summary_file(os.path.join(node_dir, "summary.csv"))
     best_results = summary_df[summary_df['filename'].isin(filenames)]
     module_params = best_results['module_params'].tolist()
-    return module_params
+    return tuple(module_params)
 
 
 def edit_summary_df_params(summary_df: pd.DataFrame, target_modules, target_module_params) -> pd.DataFrame:

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/docs/source/conf.py

@@ -10,7 +10,7 @@ project = 'AutoRAG'
 copyright = '2024, Marker-Inc'
 author = 'Marker-Inc'
 
-with open('../../VERSION') as f:
+with open('../../autorag/VERSION') as f:
     version = f.read().strip()
 
 # -- General configuration ---------------------------------------------------

AutoRAG-0.0.4/docs/source/data_creation/tutorial.md

@@ -0,0 +1 @@
+# Start creating your own evaluation data

AutoRAG-0.0.4/docs/source/index.rst

@@ -0,0 +1,45 @@
+.. AutoRAG documentation master file, created by
+   sphinx-quickstart on Wed Jan 17 20:55:21 2024.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+AutoRAG documentation
+===================================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Getting Started
+   :hidden:
+
+   install.md
+   tutorial.md
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Data Creation
+   :hidden:
+
+   data_creation/tutorial.md
+
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Optimization
+   :hidden:
+
+   optimization/optimization.md
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Available Modules
+   :hidden:
+
+   modules/index.md
+
+.. toctree::
+   :maxdepth: 1
+   :caption: API Reference
+   :hidden:
+
+   api_spec/modules
+

AutoRAG-0.0.4/docs/source/install.md

@@ -0,0 +1 @@
+# Installation

AutoRAG-0.0.4/docs/source/modules/index.md

@@ -0,0 +1 @@
+# Available Module List

AutoRAG-0.0.4/docs/source/optimization/optimization.md

@@ -0,0 +1,3 @@
+# How optimization works
+
+![Sundae](../resources/samsung_sundae.jpeg)

AutoRAG-0.0.4/docs/source/tutorial.md

@@ -0,0 +1 @@
+# Tutorial

AutoRAG-0.0.4/sample_dataset/README.md

@@ -0,0 +1,25 @@
+# sample_dataset handling
+
+The sample_dataset folder does not includes a `qa.parquet`, `corpus.parquet` file that is significantly large and cannot be uploaded directly to Git due to size limitations. 
+
+To prepare and use datasets available in the sample_dataset folder, specifically `triviaqa`, `msmarco` and `eli5`, you can follow the outlined methods below.
+
+## Usage
+
+ The example provided uses `triviaqa`, but the same approach applies to `msmarco` and `eli5`.
+
+### 1. Run with a specified save path
+To execute the Python script from the terminal and save the dataset to a specified path, use the command:
+
+```bash
+python ./sample_dataset/triviaqa/load_triviaqa_dataset.py --save_path /path/to/save/dataset
+```
+This runs the `load_triviaqa_dataset.py` script located in the `./sample_dataset/triviaqa/` directory,
+using the `--save_path` argument to specify the dataset's save location.
+
+### 2. Run without specifying a save path
+If you run the script without the `--save_path` argument, the dataset will be saved to a default location, which is the directory containing the `load_triviaqa_dataset.py` file, essentially `./sample_dataset/triviaqa/`:
+```bash
+python ./sample_dataset/triviaqa/load_triviaqa_dataset.py
+```
+This behavior allows for a straightforward execution without needing to specify a path, making it convenient for quick tests or when working directly within the target directory.

AutoRAG-0.0.4/sample_dataset/eli5/load_eli5_dataset.py

@@ -0,0 +1,30 @@
+import os
+import pathlib
+import click
+
+from datasets import load_dataset
+
+
+@click.command()
+@click.option('--save_path', type=str, default=pathlib.PurePath(__file__).parent, help='Path to save sample eli5 dataset.')
+def load_eli5_dataset(save_path):
+    # set file path
+    file_path = "MarkrAI/eli5_sample_autorag"
+
+    # load dataset
+    corpus_dataset = load_dataset(file_path, "corpus")['train'].to_pandas()
+    qa_train_dataset = load_dataset(file_path, "qa")['train'].to_pandas()
+    qa_test_dataset = load_dataset(file_path, "qa")['test'].to_pandas()
+
+    # save data
+    if os.path.exists(os.path.join(save_path, "corpus.parquet")) is True:
+        raise ValueError("corpus.parquet already exists")
+    if os.path.exists(os.path.join(save_path, "qa.parquet")) is True:
+        raise ValueError("qa.parquet already exists")
+    corpus_dataset.to_parquet(os.path.join(save_path, "corpus.parquet"))
+    qa_train_dataset.to_parquet(os.path.join(save_path, "qa_train.parquet"))
+    qa_test_dataset.to_parquet(os.path.join(save_path, "qa_test.parquet"))
+
+
+if __name__ == '__main__':
+    load_eli5_dataset()

AutoRAG-0.0.4/sample_dataset/msmarco/load_msmarco_dataset.py

@@ -0,0 +1,30 @@
+import os
+import pathlib
+import click
+
+from datasets import load_dataset
+
+
+@click.command()
+@click.option('--save_path', type=str, default=pathlib.PurePath(__file__).parent, help='Path to save sample msmarco dataset.')
+def load_msmarco_dataset(save_path):
+    # set file path
+    file_path = "MarkrAI/msmarco_sample_autorag"
+
+    # load dataset
+    corpus_dataset = load_dataset(file_path, "corpus")['train'].to_pandas()
+    qa_train_dataset = load_dataset(file_path, "qa")['train'].to_pandas()
+    qa_test_dataset = load_dataset(file_path, "qa")['test'].to_pandas()
+
+    # save corpus data
+    if os.path.exists(os.path.join(save_path, "corpus.parquet")) is True:
+        raise ValueError("corpus.parquet already exists")
+    if os.path.exists(os.path.join(save_path, "qa.parquet")) is True:
+        raise ValueError("qa.parquet already exists")
+    corpus_dataset.to_parquet(os.path.join(save_path, "corpus.parquet"), index=False)
+    qa_train_dataset.to_parquet(os.path.join(save_path, "qa_train.parquet"), index=False)
+    qa_test_dataset.to_parquet(os.path.join(save_path, "qa_test.parquet"), index=False)
+
+
+if __name__ == '__main__':
+    load_msmarco_dataset()

AutoRAG-0.0.4/sample_dataset/triviaqa/load_triviaqa_dataset.py

@@ -0,0 +1,30 @@
+import os
+import pathlib
+
+import click
+from datasets import load_dataset
+
+
+@click.command()
+@click.option('--save_path', type=str, default=pathlib.PurePath(__file__).parent, help='Path to save sample triviaqa dataset.')
+def load_triviaqa_dataset(save_path):
+    # set file path
+    file_path = "MarkrAI/triviaqa_sample_autorag"
+
+    # load dataset
+    corpus_dataset = load_dataset(file_path, "corpus")['train'].to_pandas()
+    qa_train_dataset = load_dataset(file_path, "qa")['train'].to_pandas()
+    qa_test_dataset = load_dataset(file_path, "qa")['test'].to_pandas()
+
+    # save corpus data
+    if os.path.exists(os.path.join(save_path, "corpus.parquet")) is True:
+        raise ValueError("corpus.parquet already exists")
+    if os.path.exists(os.path.join(save_path, "qa.parquet")) is True:
+        raise ValueError("qa.parquet already exists")
+    corpus_dataset.to_parquet(os.path.join(save_path, "corpus.parquet"), index=False)
+    qa_train_dataset.to_parquet(os.path.join(save_path, "qa_train.parquet"), index=False)
+    qa_test_dataset.to_parquet(os.path.join(save_path, "qa_test.parquet"), index=False)
+
+
+if __name__ == '__main__':
+    load_triviaqa_dataset()

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/tests/autorag/nodes/retrieval/test_hybrid_cc.py

@@ -16,6 +16,16 @@ def test_cc_pure():
     assert result_id == ['id-1', 'id-4', 'id-2']
 
 
+def test_cc_non_overlap():
+    sample_ids = (['id-1', 'id-2', 'id-3', 'id-4', 'id-5'],
+                  ['id-6', 'id-4', 'id-3', 'id-7', 'id-2'])
+    sample_scores = ([5, 3, 1, 0.4, 0.2], [6, 2, 1, 0.5, 0.1])
+    result_id, result_scores = cc_pure(sample_ids, sample_scores,
+                                       weights=(0.3, 0.7), top_k=3)
+    assert result_scores == pytest.approx([0.7, 0.3, 0.23792372])
+    assert result_id == ['id-6', 'id-1', 'id-4']
+
+
 def test_hybrid_cc():
     sample_ids = ([
                       ['id-1', 'id-2', 'id-3', 'id-4', 'id-5'],
@@ -58,7 +68,7 @@ def test_hybrid_cc_node(pseudo_project_dir, pseudo_node_dir):
                  ['id-1', 'id-2', 'id-3', 'id-4', 'id-5']],
                 [['id-1', 'id-4', 'id-3', 'id-5', 'id-2'],
                  ['id-1', 'id-4', 'id-3', 'id-5', 'id-2']]
-        ),
+                ),
         'scores': ([[5, 3, 1, 0.4, 0.2], [5, 3, 1, 0.4, 0.2]],
                    [[6, 2, 1, 0.5, 0.1], [6, 2, 1, 0.5, 0.1]]),
         'top_k': 3,

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/tests/autorag/nodes/retrieval/test_run_retrieval_node.py

@@ -95,7 +95,7 @@ def test_run_retrieval_node(node_line_dir):
     assert all(hybrid_summary_df['module_params'].apply(lambda x: 'target_module_params' in x))
     assert all(hybrid_summary_df['module_params'].apply(lambda x: x['target_modules'] == ('bm25', 'vectordb')))
     assert all(hybrid_summary_df['module_params'].apply(
-        lambda x: x['target_module_params'] == [{'top_k': 4}, {'top_k': 4, 'embedding_model': 'openai'}]))
+        lambda x: x['target_module_params'] == ({'top_k': 4}, {'top_k': 4, 'embedding_model': 'openai'})))
 
     # test the best file is saved properly
     best_filename = summary_df[summary_df['is_best'] == True]['filename'].values[0]
@@ -181,3 +181,48 @@ def test_select_result_for_hybrid(pseudo_node_dir):
     assert scores[1] == [[0.5, 0.6, 0.7],
                          [0.5, 0.6, 0.7],
                          [0.5, 0.6, 0.7]]
+
+
+def test_run_retrieval_node_only_hybrid(node_line_dir):
+    modules = [hybrid_cc]
+    module_params = [
+        {'top_k': 4, 'target_modules': ('bm25', 'vectordb'), 'weights': (0.3, 0.7),
+         'target_module_params': ({'top_k': 3}, {'top_k': 3, 'embedding_model': 'openai'})},
+    ]
+    project_dir = pathlib.PurePath(node_line_dir).parent.parent
+    qa_path = os.path.join(project_dir, "data", "qa.parquet")
+    strategies = {
+        'metrics': ['retrieval_f1', 'retrieval_recall'],
+    }
+    previous_result = pd.read_parquet(qa_path)
+    best_result = run_retrieval_node(modules, module_params, previous_result, node_line_dir, strategies)
+    assert os.path.exists(os.path.join(node_line_dir, "retrieval"))
+    expect_columns = ['qid', 'query', 'retrieval_gt', 'generation_gt',
+                      'retrieved_contents', 'retrieved_ids', 'retrieve_scores', 'retrieval_f1', 'retrieval_recall']
+    assert all([expect_column in best_result.columns for expect_column in expect_columns])
+    # test summary feature
+    summary_path = os.path.join(node_line_dir, "retrieval", "summary.csv")
+    single_result_path = os.path.join(node_line_dir, "retrieval", "0.parquet")
+    assert os.path.exists(single_result_path)
+    single_result_df = pd.read_parquet(single_result_path)
+    assert os.path.exists(summary_path)
+    summary_df = load_summary_file(summary_path)
+    assert set(summary_df.columns) == {'filename', 'retrieval_f1', 'retrieval_recall',
+                                       'module_name', 'module_params', 'execution_time', 'is_best'}
+    assert len(summary_df) == 1
+    assert summary_df['filename'][0] == "0.parquet"
+    assert summary_df['retrieval_f1'][0] == single_result_df['retrieval_f1'].mean()
+    assert summary_df['retrieval_recall'][0] == single_result_df['retrieval_recall'].mean()
+    assert summary_df['module_name'][0] == "hybrid_cc"
+    assert summary_df['module_params'][0] == {'top_k': 4, 'target_modules': ('bm25', 'vectordb'), 'weights': (0.3, 0.7),
+         'target_module_params': ({'top_k': 3}, {'top_k': 3, 'embedding_model': 'openai'})}
+    assert summary_df['execution_time'][0] > 0
+    assert summary_df['is_best'][0] == True
+    assert summary_df['filename'].nunique() == len(summary_df)
+
+    # test the best file is saved properly
+    best_filename = summary_df[summary_df['is_best'] == True]['filename'].values[0]
+    best_path = os.path.join(node_line_dir, "retrieval", f'best_{best_filename}')
+    assert os.path.exists(best_path)
+    best_df = pd.read_parquet(best_path)
+    assert all([expect_column in best_df.columns for expect_column in expect_columns])

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/tests/autorag/test_deploy.py

@@ -9,7 +9,8 @@ import yaml
 
 from click.testing import CliRunner
 from fastapi.testclient import TestClient
-from autorag.deploy import summary_df_to_yaml, extract_best_config, Runner, extract_node_line_names
+from autorag.deploy import summary_df_to_yaml, extract_best_config, Runner, extract_node_line_names, \
+    extract_node_strategy
 from autorag.evaluator import Evaluator, cli
 
 root_dir = pathlib.PurePath(os.path.dirname(os.path.realpath(__file__))).parent
@@ -58,6 +59,9 @@ solution_dict = {
             'nodes': [
                 {
                     'node_type': 'retrieval',
+                    'strategy': {
+                        'metrics': ['retrieval_f1', 'retrieval_recall', 'retrieval_precision'],
+                    },
                     'modules': [
                         {
                             'module_type': 'bm25',
@@ -67,6 +71,10 @@ solution_dict = {
                 },
                 {
                     'node_type': 'rerank',
+                    'strategy': {
+                        'metrics': ['retrieval_f1', 'retrieval_recall', 'retrieval_precision'],
+                        'speed_threshold': 10,
+                    },
                     'modules': [
                         {
                             'module_type': 'upr',
@@ -82,6 +90,9 @@ solution_dict = {
             'nodes': [
                 {
                     'node_type': 'generation',
+                    'strategy': {
+                        'metrics': ['bleu', 'rouge'],
+                    },
                     'modules': [
                         {
                             'module_type': 'gpt-4',
@@ -111,6 +122,18 @@ def test_extract_node_line_names(full_config):
     assert node_line_names == ['pre_retrieve_node_line', 'retrieve_node_line', 'post_retrieve_node_line']
 
 
+def test_extract_node_strategy(full_config):
+    node_strategies = extract_node_strategy(full_config)
+    assert set(list(node_strategies.keys())) == {
+        'query_expansion', 'retrieval', 'passage_reranker', 'passage_compressor',
+        'prompt_maker', 'generator'
+    }
+    assert node_strategies['retrieval'] == {
+        'metrics': ['retrieval_f1', 'retrieval_recall', 'retrieval_precision'],
+        'speed_threshold': 10,
+    }
+
+
 def test_summary_df_to_yaml():
     yaml_dict = summary_df_to_yaml(summary_df, solution_dict)
     assert yaml_dict == solution_dict

{AutoRAG-0.0.3 → AutoRAG-0.0.4}/tests/autorag/test_evaluator.py

@@ -2,10 +2,12 @@ import os.path
 import pathlib
 import shutil
 import subprocess
+import tempfile
 
 import pandas as pd
 import pytest
 
+from autorag.deploy import extract_best_config
 from autorag.evaluator import Evaluator
 from autorag.nodes.retrieval import bm25, vectordb, hybrid_rrf
 from autorag.nodes.retrieval.run import run_retrieval_node
@@ -35,6 +37,24 @@ def evaluator():
             pass
 
 
+@pytest.fixture
+def test_evaluator():
+    evaluator = Evaluator(os.path.join(resource_dir, 'qa_test_data_sample.parquet'),
+                          os.path.join(resource_dir, 'corpus_data_sample.parquet'))
+    yield evaluator
+    paths_to_remove = ['0', 'data', 'resources', 'trial.json']
+
+    for path in paths_to_remove:
+        full_path = os.path.join(os.getcwd(), path)
+        try:
+            if os.path.isdir(full_path):
+                shutil.rmtree(full_path)
+            else:
+                os.remove(full_path)
+        except FileNotFoundError:
+            pass
+
+
 def test_evaluator_init(evaluator):
     validate_qa_dataset(evaluator.qa_data)
     validate_corpus_dataset(evaluator.corpus_data)
@@ -98,7 +118,7 @@ def test_start_trial(evaluator):
     # test node line summary
     node_line_summary_path = os.path.join(os.getcwd(), '0', 'retrieve_node_line', 'summary.csv')
     assert os.path.exists(node_line_summary_path)
-    node_line_summary_df = load_summary_file(node_line_summary_path,["best_module_params"])
+    node_line_summary_df = load_summary_file(node_line_summary_path, ["best_module_params"])
     assert len(node_line_summary_df) == 1
     assert set(node_line_summary_df.columns) == {'node_type', 'best_module_filename',
                                                  'best_module_name', 'best_module_params', 'best_execution_time'}
@@ -180,3 +200,28 @@ def start_trial_full(evaluator):
     assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line', 'generator', '3.parquet'))
     assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line', 'generator', '4.parquet'))
     assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line', 'generator', '5.parquet'))
+
+
+def test_test_data_evaluate(test_evaluator):
+    trial_folder = os.path.join(resource_dir, 'result_project', '0')
+    with tempfile.NamedTemporaryFile(mode="w+t", suffix=".yaml") as yaml_file:
+        extract_best_config(trial_folder, yaml_file.name)
+        test_evaluator.start_trial(yaml_file.name)
+
+    assert os.path.exists(os.path.join(os.getcwd(), '0'))
+    assert os.path.exists(os.path.join(os.getcwd(), 'data'))
+    assert os.path.exists(os.path.join(os.getcwd(), 'resources'))
+    assert os.path.exists(os.path.join(os.getcwd(), 'trial.json'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'config.yaml'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'retrieve_node_line'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'retrieve_node_line', 'retrieval'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'retrieve_node_line', 'retrieval', '0.parquet'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'retrieve_node_line', 'retrieval', 'best_0.parquet'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'retrieve_node_line', 'summary.csv'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'pre_retrieve_node_line'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'pre_retrieve_node_line', 'query_expansion'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'pre_retrieve_node_line', 'query_expansion', "0.parquet"))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line', 'prompt_maker'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'post_retrieve_node_line', 'generator'))
+    assert os.path.exists(os.path.join(os.getcwd(), '0', 'summary.csv'))

AutoRAG-0.0.3/autorag/VERSION

@@ -1 +0,0 @@
-0.0.3