AutoRAG 0.2.3__tar.gz → 0.2.4__tar.gz

{autorag-0.2.3 → autorag-0.2.4}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.3
+Version: 0.2.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

{autorag-0.2.3 → autorag-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.3
+Version: 0.2.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

autorag-0.2.4/autorag/VERSION

@@ -0,0 +1 @@
+0.2.4

{autorag-0.2.3 → autorag-0.2.4}/autorag/dashboard.py

@@ -8,11 +8,13 @@ import pandas as pd
 import panel as pn
 import seaborn as sns
 import yaml
+from bokeh.models import NumberFormatter, BooleanFormatter
 
 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")
+             console_output='disable', sizing_mode="stretch_width",
+             css_files=["https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/all.min.css"])
 logger = logging.getLogger("AutoRAG")
 
 
@@ -71,9 +73,26 @@ 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
+    bokeh_formatters = {
+        'float': NumberFormatter(format='0.000'),
+        'bool': BooleanFormatter(),
+    }
+    first_df = pd.read_parquet(os.path.join(node_dir, '0.parquet'))
+
+    each_module_df_widget = pn.widgets.Tabulator(pd.DataFrame(columns=first_df.columns), name='Module DataFrame',
+                                                 formatters=bokeh_formatters,
+                                                 pagination='local', page_size=20, widths=150)
+
+    def change_module_widget(event):
+        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_widget.value = each_module_df
+
+    df_widget = pn.widgets.Tabulator(summary_df, name='Summary DataFrame', formatters=bokeh_formatters,
+                                     buttons={'detail': '<i class="fa fa-eye"></i>'}, widths=150)
+    df_widget.on_click(change_module_widget)
 
     try:
         fig, ax = plt.subplots(figsize=(10, 5))
@@ -86,7 +105,8 @@ def node_view(node_dir: str):
         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)
+        layout = pn.Column("## Summary distribution plot", plot_pane, "## Summary DataFrame", df_widget,
+                           "## Module Result DataFrame", each_module_df_widget)
     except Exception as e:
         logger.error(f'Skipping make boxplot and stripplot with error {e}')
         layout = pn.Column("## Summary DataFrame", df_widget)

{autorag-0.2.3 → autorag-0.2.4}/autorag/evaluator.py

@@ -121,6 +121,14 @@ class Evaluator:
             embedding_models_list = list(chain.from_iterable(
                 map(lambda nodes: extract_values_from_nodes(nodes, 'embedding_model'), node_lines.values())))
 
+            # get embedding batch size in nodes
+            embedding_batch_list = list(chain.from_iterable(
+                map(lambda nodes: extract_values_from_nodes(nodes, 'embedding_batch'), node_lines.values())))
+            if len(embedding_batch_list) == 0:
+                embedding_batch = 100
+            else:
+                embedding_batch = embedding_batch_list[0]
+
             # duplicate check in embedding_models
             embedding_models_list = list(set(embedding_models_list))
 
@@ -140,7 +148,7 @@ class Evaluator:
                 else:
                     logger.error(f"embedding_model_str {embedding_model_str} does not exist.")
                     raise KeyError(f"embedding_model_str {embedding_model_str} does not exist.")
-                vectordb_ingest(collection, self.corpus_data, embedding_model)
+                vectordb_ingest(collection, self.corpus_data, embedding_model, embedding_batch=embedding_batch)
                 logger.info(f'VectorDB corpus embedding complete with {embedding_model_str}.')
                 del embedding_model
                 if torch.cuda.is_available():

{autorag-0.2.3 → autorag-0.2.4}/autorag/nodes/retrieval/vectordb.py

@@ -15,7 +15,7 @@ from autorag.utils.util import process_batch, openai_truncate_by_token
 @retrieval_node
 def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collection,
              embedding_model: BaseEmbedding,
-             batch: int = 128) -> Tuple[List[List[str]], List[List[float]]]:
+             embedding_batch: int = 128) -> Tuple[List[List[str]], List[List[float]]]:
     """
     VectorDB retrieval function.
     You have to get a chroma collection that is already ingested.
@@ -26,7 +26,7 @@ def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collecti
     :param top_k: The number of passages to be retrieved.
     :param collection: A chroma collection instance that will be used to retrieve passages.
     :param embedding_model: An embedding model instance that will be used to embed queries.
-    :param batch: The number of queries to be processed in parallel.
+    :param embedding_batch: The number of queries to be processed in parallel.
         This is used to prevent API error at the query embedding.
         Default is 128.
 
@@ -39,7 +39,7 @@ def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collecti
     # run async vector_db_pure function
     tasks = [vectordb_pure(input_queries, top_k, collection, embedding_model) for input_queries in queries]
     loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+    results = loop.run_until_complete(process_batch(tasks, batch_size=embedding_batch))
     id_result = list(map(lambda x: x[0], results))
     score_result = list(map(lambda x: x[1], results))
     return id_result, score_result
@@ -77,8 +77,8 @@ async def vectordb_pure(queries: List[str], top_k: int, collection: chromadb.Col
 
 
 def vectordb_ingest(collection: chromadb.Collection, corpus_data: pd.DataFrame, embedding_model: BaseEmbedding,
-                    batch: int = 128):
-    embedding_model.embed_batch_size = batch
+                    embedding_batch: int = 128):
+    embedding_model.embed_batch_size = embedding_batch
     validate_corpus_dataset(corpus_data)
     ids = corpus_data['doc_id'].tolist()
 

{autorag-0.2.3 → autorag-0.2.4}/autorag/schema/node.py

@@ -75,12 +75,12 @@ def extract_values(node: Node, key: str) -> List[str]:
         if key not in module.module_param:
             return []
         value = module.module_param[key]
-        if isinstance(value, str):
+        if isinstance(value, str) or isinstance(value, int):
             return [value]
         elif isinstance(value, list):
             return value
         else:
-            raise ValueError(f"{key} must be str or list, but got {type(value)}")
+            raise ValueError(f"{key} must be str,list or int, but got {type(value)}")
 
     values = list(map(extract_module_values, node.modules))
     return list(set(list(itertools.chain.from_iterable(values))))

{autorag-0.2.3 → autorag-0.2.4}/autorag/utils/preprocess.py

@@ -44,6 +44,7 @@ def cast_qa_dataset(df: pd.DataFrame):
         else:
             raise ValueError(f"generation_gt must be str or list, but got {type(gt)}")
 
+    df = df.reset_index(drop=True)
     validate_qa_dataset(df)
     assert df['qid'].apply(lambda x: isinstance(x, str)).sum() == len(df), \
         "qid must be string type."
@@ -57,6 +58,7 @@ def cast_qa_dataset(df: pd.DataFrame):
 
 
 def cast_corpus_dataset(df: pd.DataFrame):
+    df = df.reset_index(drop=True)
     validate_corpus_dataset(df)
 
     def make_datetime_metadata(x):

{autorag-0.2.3 → autorag-0.2.4}/docs/source/nodes/retrieval/vectordb.md

@@ -18,9 +18,15 @@ Information about the Embedding model can be found [Supporting Embedding models]
 Plus, you can learn about how to add custom embedding model at [here](../../local_model.md#add-your-embedding-models). 
 ```
 
+- **Parameter**: `embedding_batch`
+- **Usage**: It is the batch size of the embedding model. It automatically set to the ingestion process using embedding
+  model.
+  If you get error on the embedding model, try to lower this parameter.
+
 ## **Example config.yaml**
 ```yaml
 modules:
   - module_type: vectordb
     embedding_model: openai
+    embedding_batch: 64
 ```

{autorag-0.2.3 → autorag-0.2.4}/sample_config/full.yaml

@@ -38,6 +38,7 @@ node_lines:
         - module_type: bm25
         - module_type: vectordb
           embedding_model: openai
+          embedding_batch: 256
         - module_type: hybrid_rrf
           target_modules: ('bm25', 'vectordb')
           rrf_k: [3, 5, 10]

{autorag-0.2.3 → autorag-0.2.4}/tests/autorag/test_dashboard.py

@@ -4,6 +4,7 @@ import pathlib
 import pandas as pd
 import pytest
 
+from autorag import dashboard
 from autorag.dashboard import get_metric_values, make_trial_summary_md
 
 root_dir = pathlib.PurePath(os.path.dirname(os.path.realpath(__file__))).parent
@@ -28,5 +29,7 @@ def test_make_trial_summary_md():
     md_text = make_trial_summary_md(sample_trial_dir)
     assert bool(md_text)
 
-# def test_dashboard_run():
-#     dashboard.run(sample_trial_dir)
+
+@pytest.mark.skip(reason="Can't stop this test on the github action or pytest cli setup")
+def test_dashboard_run():
+    dashboard.run(sample_trial_dir)

{autorag-0.2.3 → autorag-0.2.4}/tests/autorag/test_evaluator.py

@@ -69,7 +69,8 @@ def test_load_node_line(evaluator):
     assert node.modules[1].module == vectordb
     assert node.modules[2].module == hybrid_rrf
     assert node.modules[0].module_param == {'bm25_tokenizer': ['facebook/opt-125m', 'porter_stemmer']}
-    assert node.modules[1].module_param == {'embedding_model': ['openai', 'openai']}
+    assert node.modules[1].module_param == {'embedding_model': ['openai', 'openai'],
+                                            'embedding_batch': 50}
     assert node.modules[2].module_param == {
         'rrf_k': 5, 'target_modules': ('bm25', 'vectordb')
     }

{autorag-0.2.3 → autorag-0.2.4}/tests/autorag/utils/test_preprocess.py

@@ -12,9 +12,9 @@ from autorag.utils import (validate_qa_dataset, validate_corpus_dataset, cast_qa
 @pytest.fixture
 def qa_df():
     return pd.DataFrame({
-        'qid': ['id1', 'id2'],
-        'query': ['query1', 'query2'],
-        'retrieval_gt': [[['doc1', 'doc3'], ['doc2']], [[]]],
+        'qid': ['id1', 'id2', 'id3'],
+        'query': ['query1', 'query2', 'query3'],
+        'retrieval_gt': [[['doc1', 'doc3'], ['doc2']], [[]], [[]]],
         'generation_gt': 'answer1',
     })
 
@@ -55,6 +55,7 @@ def test_validate_corpus_dataset(corpus_df):
 # Test cast_qa_dataset
 def test_cast_qa_dataset(qa_df):
     # Cast the dataset and check for correct casting
+    qa_df.drop(index=1, inplace=True)
     casted_df = cast_qa_dataset(qa_df)
     assert all(isinstance(x, list) for x in casted_df['retrieval_gt'])
     assert all(isinstance(x[0], list) for x in casted_df['retrieval_gt'])
@@ -70,14 +71,12 @@ def test_cast_qa_dataset(qa_df):
 # Test cast_corpus_dataset
 def test_cast_corpus_dataset(corpus_df):
     # Cast the dataset and check for a datetime key in metadata
+    corpus_df.drop(index=1, inplace=True)
     casted_df = cast_corpus_dataset(corpus_df)
     assert all('last_modified_datetime' in x for x in casted_df['metadata'])
     assert casted_df['metadata'].iloc[0]['prev_id'] is None
     assert casted_df['metadata'].iloc[0]['next_id'] == 'doc2'
-    assert casted_df['metadata'].iloc[1]['test_key'] == 'test_value'
-    assert casted_df['metadata'].iloc[1]['prev_id'] is None
-    assert casted_df['metadata'].iloc[1]['next_id'] is None
-    assert casted_df['metadata'].iloc[2]['last_modified_datetime'] == datetime(2022, 12, 1, 3, 4, 5)
+    assert casted_df['metadata'].iloc[1]['last_modified_datetime'] == datetime(2022, 12, 1, 3, 4, 5)
 
 
 def test_validate_qa_from_corpus_dataset(qa_df, corpus_df):

{autorag-0.2.3 → autorag-0.2.4}/tests/resources/simple.yaml

@@ -11,6 +11,7 @@ node_lines:
           bm25_tokenizer: [ facebook/opt-125m, porter_stemmer ]
         - module_type: vectordb
           embedding_model: [openai, openai]
+          embedding_batch: 50
         - module_type: hybrid_rrf
           rrf_k: 5
           target_modules: ('bm25', 'vectordb') # it will select the best result among each module and fuse it.

autorag-0.2.3/autorag/VERSION

@@ -1 +0,0 @@
-0.2.3