AutoRAG 0.1.12__tar.gz → 0.2.0__tar.gz

{autorag-0.1.12 → autorag-0.2.0}/.github/workflows/test.yml

@@ -7,9 +7,6 @@ on:
   pull_request:
     branches:
       - main
-  pull_request_target:
-    branches:
-      - main
 
 env:
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}

{autorag-0.1.12 → autorag-0.2.0}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.1.12
+Version: 0.2.0
 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
@@ -260,6 +260,11 @@ Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-retrievers-bm25
 Requires-Dist: streamlit
 Requires-Dist: langchain-core>=0.1.6
+Requires-Dist: panel
+Requires-Dist: seaborn
+Requires-Dist: ipykernel
+Requires-Dist: ipywidgets
+Requires-Dist: ipywidgets_bokeh
 
 # AutoRAG
 
@@ -295,7 +300,8 @@ You can see on [YouTube](https://youtu.be/2ojK8xjyXAU?feature=shared)
   - [1. Prepare your evaluation data](#1-prepare-your-evaluation-data)
   - [2. Evaluate your data to various RAG modules](#2-evaluate-your-data-to-various-rag-modules)
   - [3. Use a found optimal RAG pipeline](#3-use-a-found-optimal-rag-pipeline)
-  - [4. Share your RAG pipeline](#4-share-your-rag-pipeline)
+  - [4. Run Dashboard to see the result](#4-run-dashboard)
+  - [5. Share your RAG pipeline](#5-share-your-rag-pipeline)
   - [+ Config yaml file](#-create-your-own-config-yaml-file)
 - [Supporting RAG modules](#supporting-nodes--modules)
 - [Roadmap](#roadmap)
@@ -425,7 +431,15 @@ You can run api server with CLI command.
 autorag run_api --config_path your/path/to/pipeline.yaml --host 0.0.0.0 --port 8000
 ```
 
-### 4. Share your RAG pipeline
+### 4. Run Dashboard
+
+You can run dashboard to easily see the result.
+
+```bash
+autorag dashboard --trial_dir /your/path/to/trial_dir
+```
+
+### 5. Share your RAG pipeline
 
 You can use your RAG pipeline from extracted pipeline yaml file.
 This extracted pipeline is great for sharing your RAG pipeline to others.
@@ -515,12 +529,12 @@ at [here](https://edai.notion.site/Supporting-metrics-867d71caefd7401c9264dd91ba
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline
-- [ ] Visualize evaluation result
+- [x] Visualize evaluation result
 - [ ] Visualize config yaml file
-- [ ] More RAG modules support
+- [x] More RAG modules support
 - [x] Token usage strategy
 - [ ] Multi-modal support
-- [ ] More evaluation metrics
+- [x] More evaluation metrics
 - [ ] Answer Filtering Module
 - [x] Restart optimization from previous trial
 

{autorag-0.1.12 → autorag-0.2.0}/AutoRAG.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ AutoRAG.egg-info/top_level.txt
 autorag/VERSION
 autorag/__init__.py
 autorag/cli.py
+autorag/dashboard.py
 autorag/deploy.py
 autorag/evaluator.py
 autorag/node_line.py
@@ -45,7 +46,6 @@ autorag/evaluate/metric/__init__.py
 autorag/evaluate/metric/generation.py
 autorag/evaluate/metric/retrieval.py
 autorag/evaluate/metric/retrieval_contents.py
-autorag/evaluate/metric/retrieval_no_gt.py
 autorag/evaluate/metric/util.py
 autorag/evaluate/metric/g_eval_prompts/coh_detailed.txt
 autorag/evaluate/metric/g_eval_prompts/con_detailed.txt
@@ -234,6 +234,7 @@ docs/source/optimization/custom_config.md
 docs/source/optimization/folder_structure.md
 docs/source/optimization/optimization.md
 docs/source/optimization/sample_full_config.yaml
+docs/source/optimization/strategies.md
 docs/source/roadmap/modular_rag.md
 sample_config/compact_local.yaml
 sample_config/compact_openai.yaml
@@ -251,6 +252,7 @@ tests/delete_tests.py
 tests/mock.py
 tests/requirements.txt
 tests/autorag/test_cli.py
+tests/autorag/test_dashboard.py
 tests/autorag/test_deploy.py
 tests/autorag/test_evaluator.py
 tests/autorag/test_strategy.py
@@ -270,7 +272,6 @@ tests/autorag/evaluate/test_retrieval_evaluate.py
 tests/autorag/evaluate/metric/test_generation_metric.py
 tests/autorag/evaluate/metric/test_retrieval_contents_metric.py
 tests/autorag/evaluate/metric/test_retrieval_metric.py
-tests/autorag/evaluate/metric/test_retrieval_no_gt_metric.py
 tests/autorag/nodes/generator/test_generator_base.py
 tests/autorag/nodes/generator/test_llama_index_llm.py
 tests/autorag/nodes/generator/test_openai.py
@@ -386,7 +387,7 @@ tests/resources/result_project/1/pre_retrieve_node_line/query_expansion/1.parque
 tests/resources/result_project/1/pre_retrieve_node_line/query_expansion/2.parquet
 tests/resources/result_project/1/pre_retrieve_node_line/query_expansion/best_0.parquet
 tests/resources/result_project/1/pre_retrieve_node_line/query_expansion/summary.csv
-tests/resources/result_project/1/retrieve_node_line/passage_compressor/0.parquet
+tests/resources/result_project/1/retrieve_node_line/passage_filter/0.parquet
 tests/resources/result_project/1/retrieve_node_line/passage_reranker/0.parquet
 tests/resources/result_project/1/retrieve_node_line/passage_reranker/1.parquet
 tests/resources/result_project/1/retrieve_node_line/passage_reranker/best_0.parquet

{autorag-0.1.12 → autorag-0.2.0}/AutoRAG.egg-info/requires.txt

@@ -39,3 +39,8 @@ llama-index-llms-openai-like
 llama-index-retrievers-bm25
 streamlit
 langchain-core>=0.1.6
+panel
+seaborn
+ipykernel
+ipywidgets
+ipywidgets_bokeh

{autorag-0.1.12 → autorag-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.1.12
+Version: 0.2.0
 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
@@ -260,6 +260,11 @@ Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-retrievers-bm25
 Requires-Dist: streamlit
 Requires-Dist: langchain-core>=0.1.6
+Requires-Dist: panel
+Requires-Dist: seaborn
+Requires-Dist: ipykernel
+Requires-Dist: ipywidgets
+Requires-Dist: ipywidgets_bokeh
 
 # AutoRAG
 
@@ -295,7 +300,8 @@ You can see on [YouTube](https://youtu.be/2ojK8xjyXAU?feature=shared)
   - [1. Prepare your evaluation data](#1-prepare-your-evaluation-data)
   - [2. Evaluate your data to various RAG modules](#2-evaluate-your-data-to-various-rag-modules)
   - [3. Use a found optimal RAG pipeline](#3-use-a-found-optimal-rag-pipeline)
-  - [4. Share your RAG pipeline](#4-share-your-rag-pipeline)
+  - [4. Run Dashboard to see the result](#4-run-dashboard)
+  - [5. Share your RAG pipeline](#5-share-your-rag-pipeline)
   - [+ Config yaml file](#-create-your-own-config-yaml-file)
 - [Supporting RAG modules](#supporting-nodes--modules)
 - [Roadmap](#roadmap)
@@ -425,7 +431,15 @@ You can run api server with CLI command.
 autorag run_api --config_path your/path/to/pipeline.yaml --host 0.0.0.0 --port 8000
 ```
 
-### 4. Share your RAG pipeline
+### 4. Run Dashboard
+
+You can run dashboard to easily see the result.
+
+```bash
+autorag dashboard --trial_dir /your/path/to/trial_dir
+```
+
+### 5. Share your RAG pipeline
 
 You can use your RAG pipeline from extracted pipeline yaml file.
 This extracted pipeline is great for sharing your RAG pipeline to others.
@@ -515,12 +529,12 @@ at [here](https://edai.notion.site/Supporting-metrics-867d71caefd7401c9264dd91ba
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline
-- [ ] Visualize evaluation result
+- [x] Visualize evaluation result
 - [ ] Visualize config yaml file
-- [ ] More RAG modules support
+- [x] More RAG modules support
 - [x] Token usage strategy
 - [ ] Multi-modal support
-- [ ] More evaluation metrics
+- [x] More evaluation metrics
 - [ ] Answer Filtering Module
 - [x] Restart optimization from previous trial
 

{autorag-0.1.12 → autorag-0.2.0}/README.md

@@ -32,7 +32,8 @@ You can see on [YouTube](https://youtu.be/2ojK8xjyXAU?feature=shared)
   - [1. Prepare your evaluation data](#1-prepare-your-evaluation-data)
   - [2. Evaluate your data to various RAG modules](#2-evaluate-your-data-to-various-rag-modules)
   - [3. Use a found optimal RAG pipeline](#3-use-a-found-optimal-rag-pipeline)
-  - [4. Share your RAG pipeline](#4-share-your-rag-pipeline)
+  - [4. Run Dashboard to see the result](#4-run-dashboard)
+  - [5. Share your RAG pipeline](#5-share-your-rag-pipeline)
   - [+ Config yaml file](#-create-your-own-config-yaml-file)
 - [Supporting RAG modules](#supporting-nodes--modules)
 - [Roadmap](#roadmap)
@@ -162,7 +163,15 @@ You can run api server with CLI command.
 autorag run_api --config_path your/path/to/pipeline.yaml --host 0.0.0.0 --port 8000
 ```
 
-### 4. Share your RAG pipeline
+### 4. Run Dashboard
+
+You can run dashboard to easily see the result.
+
+```bash
+autorag dashboard --trial_dir /your/path/to/trial_dir
+```
+
+### 5. Share your RAG pipeline
 
 You can use your RAG pipeline from extracted pipeline yaml file.
 This extracted pipeline is great for sharing your RAG pipeline to others.
@@ -252,12 +261,12 @@ at [here](https://edai.notion.site/Supporting-metrics-867d71caefd7401c9264dd91ba
 # 🛣Roadmap
 
 - [ ] Policy Module for modular RAG pipeline
-- [ ] Visualize evaluation result
+- [x] Visualize evaluation result
 - [ ] Visualize config yaml file
-- [ ] More RAG modules support
+- [x] More RAG modules support
 - [x] Token usage strategy
 - [ ] Multi-modal support
-- [ ] More evaluation metrics
+- [x] More evaluation metrics
 - [ ] Answer Filtering Module
 - [x] Restart optimization from previous trial
 

autorag-0.2.0/autorag/VERSION

@@ -0,0 +1 @@
+0.2.0

{autorag-0.1.12 → autorag-0.2.0}/autorag/__init__.py

@@ -4,6 +4,7 @@ import os
 import sys
 
 import transformers
+from llama_index.core.llms.mock import MockLLM
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
 from llama_index.embeddings.openai import OpenAIEmbeddingModelType
@@ -55,6 +56,7 @@ generator_models = {
     'openai': OpenAI,
     'huggingfacellm': HuggingFaceLLM,
     'openailike': OpenAILike,
+    'mock': MockLLM,
 }
 
 rich_format = "[%(filename)s:%(lineno)s] >> %(message)s"

{autorag-0.1.12 → autorag-0.2.0}/autorag/cli.py

@@ -8,6 +8,7 @@ from typing import Optional
 
 import click
 
+from autorag import dashboard
 from autorag.deploy import Runner
 from autorag.deploy import extract_best_config as original_extract_best_config
 from autorag.evaluator import Evaluator
@@ -70,6 +71,12 @@ def run_web(yaml_path: Optional[str], project_dir: Optional[str], trial_path: Op
         subprocess.run(['streamlit', 'run', web_py_path, '--', '--trial_path', trial_path])
 
 
+@click.command()
+@click.option('--trial_dir', type=click.Path(dir_okay=True, file_okay=False, exists=True), required=True)
+def run_dashboard(trial_dir: str):
+    dashboard.run(trial_dir)
+
+
 @click.command()
 @click.option('--trial_path', type=click.Path(), help='Path to the trial directory.')
 @click.option('--output_path', type=click.Path(), help='Path to the output directory.'
@@ -94,6 +101,7 @@ def restart_evaluate(trial_path):
 cli.add_command(evaluate, 'evaluate')
 cli.add_command(run_api, 'run_api')
 cli.add_command(run_web, 'run_web')
+cli.add_command(run_dashboard, 'dashboard')
 cli.add_command(extract_best_config, 'extract_best_config')
 cli.add_command(restart_evaluate, 'restart_evaluate')
 

autorag-0.2.0/autorag/dashboard.py

@@ -0,0 +1,132 @@
+import ast
+import logging
+import os
+from typing import Dict, List
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import panel as pn
+import seaborn as sns
+import yaml
+
+from autorag.utils.util import dict_to_markdown, dict_to_markdown_table
+
+pn.extension('terminal', 'tabulator', 'mathjax', 'ipywidgets',
+             console_output='disable', sizing_mode="stretch_width")
+logger = logging.getLogger("AutoRAG")
+
+
+def find_node_dir(trial_dir: str) -> List[str]:
+    trial_summary_df = pd.read_csv(os.path.join(trial_dir, 'summary.csv'))
+    result_paths = []
+    for idx, row in trial_summary_df.iterrows():
+        node_line_name = row['node_line_name']
+        node_type = row['node_type']
+        result_paths.append(os.path.join(trial_dir, node_line_name, node_type))
+    return result_paths
+
+
+def get_metric_values(node_summary_df: pd.DataFrame) -> Dict:
+    non_metric_column_names = ['filename', 'module_name', 'module_params', 'execution_time', 'average_output_token',
+                               'is_best']
+    best_row = node_summary_df.loc[node_summary_df['is_best']].drop(columns=non_metric_column_names, errors='ignore')
+    assert len(best_row) == 1, "The best module must be only one."
+    return best_row.iloc[0].to_dict()
+
+
+def make_trial_summary_md(trial_dir):
+    markdown_text = f"""# Trial Result Summary
+- Trial Directory : {trial_dir}
+
+"""
+    node_dirs = find_node_dir(trial_dir)
+    for node_dir in node_dirs:
+        node_summary_filepath = os.path.join(node_dir, 'summary.csv')
+        node_type = os.path.basename(node_dir)
+        node_summary_df = pd.read_csv(node_summary_filepath)
+        best_row = node_summary_df.loc[node_summary_df['is_best']].iloc[0]
+        metric_dict = get_metric_values(node_summary_df)
+        markdown_text += f"""---
+
+## {node_type} best module
+
+### Module Name
+
+{best_row['module_name']}
+
+### Module Params
+
+{dict_to_markdown(ast.literal_eval(best_row['module_params']), level=3)}
+
+### Metric Values
+
+{dict_to_markdown_table(metric_dict, key_column_name='metric_name', value_column_name='metric_value')}
+
+"""
+
+    return markdown_text
+
+
+def node_view(node_dir: str):
+    non_metric_column_names = ['filename', 'module_name', 'module_params', 'execution_time', 'average_output_token',
+                               'is_best']
+    summary_df = pd.read_csv(os.path.join(node_dir, 'summary.csv'))
+    df_widget = pn.widgets.Tabulator(summary_df, name='Summary DataFrame')
+    # TODO: add on click listener for pop-up of each file detail.
+    # https://panel.holoviz.org/reference/widgets/Tabulator.html
+
+    try:
+        fig, ax = plt.subplots(figsize=(10, 5))
+        metric_df = summary_df.drop(columns=non_metric_column_names, errors='ignore')
+        sns.stripplot(data=metric_df, ax=ax)
+        strip_plot_pane = pn.pane.Matplotlib(fig, tight=True)
+
+        fig2, ax2 = plt.subplots(figsize=(10, 5))
+        sns.boxplot(data=metric_df, ax=ax2)
+        box_plot_pane = pn.pane.Matplotlib(fig2, tight=True)
+        plot_pane = pn.Row(strip_plot_pane, box_plot_pane)
+
+        layout = pn.Column("## Summary distribution plot", plot_pane, "## Summary DataFrame", df_widget)
+    except Exception as e:
+        logger.error(f'Skipping make boxplot and stripplot with error {e}')
+        layout = pn.Column("## Summary DataFrame", df_widget)
+    layout.servable()
+    return layout
+
+
+CSS = """
+div.card-margin:nth-child(1) {
+    max-height: 300px;
+}
+div.card-margin:nth-child(2) {
+    max-height: 400px;
+}
+"""
+
+
+def yaml_to_markdown(yaml_filepath):
+    markdown_content = ""
+    with open(yaml_filepath, 'r', encoding='utf-8') as file:
+        try:
+            content = yaml.safe_load(file)
+            markdown_content += f"## {os.path.basename(yaml_filepath)}\n```yaml\n{yaml.dump(content, allow_unicode=True)}\n```\n\n"
+        except yaml.YAMLError as exc:
+            print(f"Error in {yaml_filepath}: {exc}")
+    return markdown_content
+
+
+def run(trial_dir: str):
+    trial_summary_md = make_trial_summary_md(trial_dir=trial_dir)
+    trial_summary_tab = pn.pane.Markdown(trial_summary_md, sizing_mode='stretch_width')
+
+    node_views = [(str(os.path.basename(node_dir)), node_view(node_dir)) for node_dir in find_node_dir(trial_dir)]
+
+    yaml_file_markdown = yaml_to_markdown(os.path.join(trial_dir, "config.yaml"))
+
+    yaml_file = pn.pane.Markdown(yaml_file_markdown, sizing_mode='stretch_width')
+
+    tabs = pn.Tabs(('Summary', trial_summary_tab), *node_views, ('Used YAML file', yaml_file), dynamic=True)
+
+    template = pn.template.FastListTemplate(site="AutoRAG", title="Dashboard",
+                                            main=[tabs], raw_css=[CSS]).servable()
+    template.show()

autorag-0.2.0/autorag/evaluate/retrieval.py

@@ -0,0 +1,67 @@
+import functools
+import warnings
+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
+
+RETRIEVAL_METRIC_FUNC_DICT = {func.__name__: func for func in
+                              [retrieval_recall, retrieval_precision, retrieval_f1, retrieval_ndcg, retrieval_mrr,
+                               retrieval_map]}
+RETRIEVAL_NO_GT_METRIC_FUNC_DICT = {func.__name__: func for func in []}
+
+
+def evaluate_retrieval(retrieval_gt: List[List[List[str]]], metrics: Union[List[str], List[Dict]],
+                       queries: Optional[List[str]] = None, generation_gt: Optional[List[List[str]]] = None):
+    def decorator_evaluate_retrieval(
+            func: Callable[[Any], Tuple[List[List[str]], List[List[str]], List[List[float]]]]):
+        """
+        Decorator for evaluating retrieval results.
+        You can use this decorator to any method that returns (contents, scores, ids),
+        which is the output of conventional retrieval modules.
+
+        :param func: Must return (contents, scores, ids)
+        :return: wrapper function that returns pd.DataFrame, which is the evaluation result.
+        """
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs) -> pd.DataFrame:
+            contents, pred_ids, scores = func(*args, **kwargs)
+
+            metric_scores = {}
+            metric_names, metric_params = cast_metrics(metrics)
+
+            for metric_name, metric_param in zip(metric_names, metric_params):
+                if metric_name in RETRIEVAL_METRIC_FUNC_DICT:
+                    metric_func = RETRIEVAL_METRIC_FUNC_DICT[metric_name]
+                    metric_scores[metric_name] = metric_func(retrieval_gt=retrieval_gt, pred_ids=pred_ids,
+                                                             **metric_param)
+                elif metric_name in RETRIEVAL_NO_GT_METRIC_FUNC_DICT:
+                    metric_func = RETRIEVAL_NO_GT_METRIC_FUNC_DICT[metric_name]
+                    if queries is None:
+                        raise ValueError(f"To using {metric_name}, you have to provide queries.")
+                    if generation_gt is None:
+                        raise ValueError(f"To using {metric_name}, you have to provide generation ground truth.")
+                    metric_scores[metric_name] = metric_func(queries=queries, retrieved_contents=contents,
+                                                             generation_gt=generation_gt, **metric_param)
+                else:
+                    warnings.warn(
+                        f"metric {metric_name} is not in supported metrics: {RETRIEVAL_METRIC_FUNC_DICT.keys()}"
+                        f" and {RETRIEVAL_NO_GT_METRIC_FUNC_DICT.keys()}"
+                        f"{metric_name} will be ignored.")
+
+            metric_result_df = pd.DataFrame(metric_scores)
+            execution_result_df = pd.DataFrame({
+                'retrieved_contents': contents,
+                'retrieved_ids': pred_ids,
+                'retrieve_scores': scores,
+            })
+            result_df = pd.concat([execution_result_df, metric_result_df], axis=1)
+            return result_df
+
+        return wrapper
+
+    return decorator_evaluate_retrieval

{autorag-0.1.12 → autorag-0.2.0}/autorag/nodes/generator/run.py

@@ -6,7 +6,7 @@ import pandas as pd
 
 from autorag.evaluate import evaluate_generation
 from autorag.evaluate.util import cast_metrics
-from autorag.strategy import measure_speed, filter_by_threshold, select_best_average
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
 
 
 def run_generator_node(modules: List[Callable],
@@ -72,7 +72,8 @@ def run_generator_node(modules: List[Callable],
         results, filenames = filter_by_threshold(results, average_times, strategies['speed_threshold'], filenames)
     if strategies.get('token_threshold') is not None:
         results, filenames = filter_by_threshold(results, token_usages, strategies['token_threshold'], filenames)
-    selected_result, selected_filename = select_best_average(results, metric_names, filenames)
+    selected_result, selected_filename = select_best(results, metric_names, filenames,
+                                                     strategies.get('strategy', 'mean'))
     best_result = pd.concat([previous_result, selected_result], axis=1)
 
     # add 'is_best' column at summary file

{autorag-0.1.12 → autorag-0.2.0}/autorag/nodes/passageaugmenter/run.py

@@ -6,7 +6,7 @@ from typing import List, Callable, Dict
 import pandas as pd
 
 from autorag.nodes.retrieval.run import evaluate_retrieval_node
-from autorag.strategy import measure_speed, filter_by_threshold, select_best_average
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
 
 logger = logging.getLogger("AutoRAG")
 
@@ -56,7 +56,8 @@ def run_passage_augmenter_node(modules: List[Callable],
     # filter by strategies
     if strategies.get('speed_threshold') is not None:
         results, filenames = filter_by_threshold(results, average_times, strategies['speed_threshold'], filenames)
-    selected_result, selected_filename = select_best_average(results, strategies.get('metrics'), filenames)
+    selected_result, selected_filename = select_best(results, strategies.get('metrics'), filenames,
+                                                     strategies.get('strategy', 'mean'))
     # change metric name columns to passage_augmenter_metric_name
     selected_result = selected_result.rename(columns={
         metric_name: f'passage_augmenter_{metric_name}' for metric_name in strategies['metrics']})

{autorag-0.1.12 → autorag-0.2.0}/autorag/nodes/passagecompressor/run.py

@@ -6,8 +6,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.strategy import measure_speed, filter_by_threshold, select_best_average
-from autorag.utils import validate_qa_dataset, validate_corpus_dataset
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
 from autorag.utils.util import fetch_contents
 
 
@@ -40,16 +39,21 @@ def run_passage_compressor_node(modules: List[Callable],
     if not os.path.exists(save_dir):
         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"))
+    # 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."
+
     # run modules
     results, execution_times = zip(*map(lambda task: measure_speed(
         task[0], project_dir=project_dir, previous_result=previous_result, **task[1]), zip(modules, module_params)))
     results = list(results)
     average_times = list(map(lambda x: x / len(results[0]), execution_times))
 
-    # 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"))
-    retrieval_contents_gt = make_contents_gt(qa_data, corpus_data)
+    retrieval_contents_gt = list(map(lambda x: fetch_contents(corpus_data, x), qa_data['retrieval_gt'].tolist()))
+    retrieval_contents_gt = list(map(lambda x: list(itertools.chain.from_iterable(x)), retrieval_contents_gt))
 
     # run metrics before filtering
     if strategies.get('metrics') is None:
@@ -76,7 +80,8 @@ def run_passage_compressor_node(modules: List[Callable],
     # filter by strategies
     if strategies.get('speed_threshold') is not None:
         results, filenames = filter_by_threshold(results, average_times, strategies['speed_threshold'], filenames)
-    selected_result, selected_filename = select_best_average(results, strategies.get('metrics'), filenames)
+    selected_result, selected_filename = select_best(results, strategies.get('metrics'), filenames,
+                                                     strategies.get('strategy', 'mean'))
     new_retrieved_contents = selected_result['retrieved_contents']
     previous_result['retrieved_contents'] = new_retrieved_contents
     selected_result = selected_result.drop(columns=['retrieved_contents'])
@@ -114,10 +119,3 @@ def evaluate_passage_compressor_node(result_df: pd.DataFrame,
     )), metrics))
     result_df = pd.concat([result_df, pd.DataFrame(metrics_scores)], axis=1)
     return result_df
-
-
-def make_contents_gt(qa_data: pd.DataFrame, corpus_data: pd.DataFrame) -> List[List[str]]:
-    validate_qa_dataset(qa_data)
-    validate_corpus_dataset(corpus_data)
-    retrieval_gt_ids = qa_data['retrieval_gt'].apply(lambda x: list(itertools.chain.from_iterable(x))).tolist()
-    return fetch_contents(corpus_data, retrieval_gt_ids)

{autorag-0.1.12 → autorag-0.2.0}/autorag/nodes/passagefilter/run.py

@@ -5,7 +5,7 @@ from typing import List, Callable, Dict
 import pandas as pd
 
 from autorag.nodes.retrieval.run import evaluate_retrieval_node
-from autorag.strategy import measure_speed, filter_by_threshold, select_best_average
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
 
 
 def run_passage_filter_node(modules: List[Callable],
@@ -66,7 +66,8 @@ def run_passage_filter_node(modules: List[Callable],
     # filter by strategies
     if strategies.get('speed_threshold') is not None:
         results, filenames = filter_by_threshold(results, average_times, strategies['speed_threshold'], filenames)
-    selected_result, selected_filename = select_best_average(results, strategies.get('metrics'), filenames)
+    selected_result, selected_filename = select_best(results, strategies.get('metrics'), filenames,
+                                                     strategies.get('strategy', 'mean'))
     selected_result = selected_result.rename(columns={
         metric_name: f'passage_filter_{metric_name}' for metric_name in strategies['metrics']})
     previous_result = previous_result.drop(columns=['retrieved_contents', 'retrieved_ids', 'retrieve_scores'])

{autorag-0.1.12 → autorag-0.2.0}/autorag/nodes/passagereranker/run.py

@@ -6,7 +6,7 @@ from typing import List, Callable, Dict
 import pandas as pd
 
 from autorag.nodes.retrieval.run import evaluate_retrieval_node
-from autorag.strategy import measure_speed, filter_by_threshold, select_best_average
+from autorag.strategy import measure_speed, filter_by_threshold, select_best
 
 logger = logging.getLogger("AutoRAG")
 
@@ -70,7 +70,8 @@ def run_passage_reranker_node(modules: List[Callable],
     # filter by strategies
     if strategies.get('speed_threshold') is not None:
         results, filenames = filter_by_threshold(results, average_times, strategies['speed_threshold'], filenames)
-    selected_result, selected_filename = select_best_average(results, strategies.get('metrics'), filenames)
+    selected_result, selected_filename = select_best(results, strategies.get('metrics'), filenames,
+                                                     strategies.get('strategy', 'mean'))
     # change metric name columns to passage_reranker_metric_name
     selected_result = selected_result.rename(columns={
         metric_name: f'passage_reranker_{metric_name}' for metric_name in strategies['metrics']})