AutoRAG 0.2.9__tar.gz → 0.2.11__tar.gz

{autorag-0.2.9 → autorag-0.2.11}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.9
+Version: 0.2.11
 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
@@ -256,7 +256,7 @@ 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
+Requires-Dist: llama-index-llms-openai>=0.1.26
 Requires-Dist: llama-index-llms-huggingface
 Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-llms-ollama

{autorag-0.2.9 → autorag-0.2.11}/AutoRAG.egg-info/SOURCES.txt

@@ -116,9 +116,7 @@ autorag/nodes/retrieval/__init__.py
 autorag/nodes/retrieval/base.py
 autorag/nodes/retrieval/bm25.py
 autorag/nodes/retrieval/hybrid_cc.py
-autorag/nodes/retrieval/hybrid_dbsf.py
 autorag/nodes/retrieval/hybrid_rrf.py
-autorag/nodes/retrieval/hybrid_rsf.py
 autorag/nodes/retrieval/run.py
 autorag/nodes/retrieval/vectordb.py
 autorag/schema/__init__.py
@@ -235,15 +233,12 @@ docs/source/nodes/query_expansion/query_decompose.md
 docs/source/nodes/query_expansion/query_expansion.md
 docs/source/nodes/retrieval/bm25.md
 docs/source/nodes/retrieval/hybrid_cc.md
-docs/source/nodes/retrieval/hybrid_dbsf.md
 docs/source/nodes/retrieval/hybrid_rrf.md
-docs/source/nodes/retrieval/hybrid_rsf.md
 docs/source/nodes/retrieval/retrieval.md
 docs/source/nodes/retrieval/vectordb.md
 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
@@ -336,9 +331,7 @@ tests/autorag/nodes/queryexpansion/test_query_expansion_run.py
 tests/autorag/nodes/retrieval/test_bm25.py
 tests/autorag/nodes/retrieval/test_hybrid_base.py
 tests/autorag/nodes/retrieval/test_hybrid_cc.py
-tests/autorag/nodes/retrieval/test_hybrid_dbsf.py
 tests/autorag/nodes/retrieval/test_hybrid_rrf.py
-tests/autorag/nodes/retrieval/test_hybrid_rsf.py
 tests/autorag/nodes/retrieval/test_retrieval_base.py
 tests/autorag/nodes/retrieval/test_run_retrieval_node.py
 tests/autorag/nodes/retrieval/test_vectordb.py

{autorag-0.2.9 → autorag-0.2.11}/AutoRAG.egg-info/requires.txt

@@ -34,7 +34,7 @@ llama-index-core>=0.10.1
 llama-index-readers-file
 llama-index-embeddings-openai
 llama-index-embeddings-huggingface
-llama-index-llms-openai
+llama-index-llms-openai>=0.1.26
 llama-index-llms-huggingface
 llama-index-llms-openai-like
 llama-index-llms-ollama

{autorag-0.2.9 → autorag-0.2.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.9
+Version: 0.2.11
 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
@@ -256,7 +256,7 @@ 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
+Requires-Dist: llama-index-llms-openai>=0.1.26
 Requires-Dist: llama-index-llms-huggingface
 Requires-Dist: llama-index-llms-openai-like
 Requires-Dist: llama-index-llms-ollama

autorag-0.2.11/autorag/VERSION

@@ -0,0 +1 @@
+0.2.11

{autorag-0.2.9 → autorag-0.2.11}/autorag/nodes/generator/openai_llm.py

@@ -13,6 +13,8 @@ from autorag.utils.util import process_batch
 logger = logging.getLogger("AutoRAG")
 
 MAX_TOKEN_DICT = {  # model name : token limit
+    'gpt-4o-mini': 128_000,
+    'gpt-4o-mini-2024-07-18': 128_000,
     'gpt-4o': 128_000,
     'gpt-4o-2024-05-13': 128_000,
     'gpt-4-turbo': 128_000,

{autorag-0.2.9 → autorag-0.2.11}/autorag/nodes/retrieval/__init__.py

@@ -1,7 +1,5 @@
 from .base import retrieval_node
 from .bm25 import bm25
 from .hybrid_cc import hybrid_cc
-from .hybrid_dbsf import hybrid_dbsf
 from .hybrid_rrf import hybrid_rrf
-from .hybrid_rsf import hybrid_rsf
 from .vectordb import vectordb

{autorag-0.2.9 → autorag-0.2.11}/autorag/nodes/retrieval/base.py

@@ -32,7 +32,6 @@ def retrieval_node(func):
             project_dir: Union[str, Path],
             previous_result: pd.DataFrame,
             **kwargs) -> Tuple[List[List[str]], List[List[str]], List[List[float]]]:
-        logger.info(f"Running retrieval node - {func.__name__} module...")
         validate_qa_dataset(previous_result)
         resources_dir = os.path.join(project_dir, "resources")
         data_dir = os.path.join(project_dir, "data")
@@ -75,10 +74,10 @@ def retrieval_node(func):
             del embedding_model
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
-        elif func.__name__ in ["hybrid_rrf", "hybrid_cc", "hybrid_rsf", "hybrid_dbsf"]:
-            if 'ids' in kwargs and 'scores' in kwargs:
+        elif func.__name__ in ["hybrid_rrf", "hybrid_cc"]:
+            if 'ids' in kwargs and 'scores' in kwargs:  # ordinary run_evaluate
                 ids, scores = func(**kwargs)
-            else:
+            else:  # => for Runner.run
                 if not ('target_modules' in kwargs and 'target_module_params' in kwargs):
                     raise ValueError(
                         f"If there are no ids and scores specified, target_modules and target_module_params must be specified for using {func.__name__}.")

{autorag-0.2.9 → autorag-0.2.11}/autorag/nodes/retrieval/bm25.py

@@ -2,7 +2,7 @@ import asyncio
 import os
 import pickle
 import re
-from typing import List, Dict, Tuple, Callable, Union, Iterable
+from typing import List, Dict, Tuple, Callable, Union, Iterable, Optional
 
 import numpy as np
 import pandas as pd
@@ -51,7 +51,8 @@ BM25_TOKENIZER = {
 
 
 @retrieval_node
-def bm25(queries: List[List[str]], top_k: int, bm25_corpus: Dict, bm25_tokenizer: str = 'porter_stemmer') -> \
+def bm25(queries: List[List[str]], top_k: int, bm25_corpus: Dict, bm25_tokenizer: str = 'porter_stemmer',
+         ids: Optional[List[List[str]]] = None) -> \
         Tuple[List[List[str]], List[List[float]]]:
     """
     BM25 retrieval function.
@@ -74,6 +75,9 @@ def bm25(queries: List[List[str]], top_k: int, bm25_corpus: Dict, bm25_tokenizer
         It supports 'porter_stemmer', 'ko_kiwi', and huggingface `AutoTokenizer`.
         You can pass huggingface tokenizer name.
         Default is porter_stemmer.
+    :param ids: The optional list of ids that you want to retrieve.
+        You don't need to specify this in the general use cases.
+        Default is None.
     :return: The 2-d list contains a list of passage ids that retrieved from bm25 and 2-d list of its scores.
         It will be a length of queries. And each element has a length of top_k.
     """
@@ -85,6 +89,12 @@ def bm25(queries: List[List[str]], top_k: int, bm25_corpus: Dict, bm25_tokenizer
         (f"The bm25 corpus tokenizer is {bm25_corpus['tokenizer_name']}, but your input is {bm25_tokenizer}. "
          f"You need to ingest again. Delete bm25 pkl file and re-ingest it.")
     bm25_instance = BM25Okapi(bm25_corpus["tokens"])
+
+    if ids is not None:
+        score_result = list(map(lambda query_list, id_list: get_bm25_scores(
+            query_list, id_list, tokenizer, bm25_instance, bm25_corpus), queries, ids))
+        return ids, score_result
+
     # run async bm25_pure function
     tasks = [bm25_pure(input_queries, top_k, tokenizer, bm25_instance, bm25_corpus) for input_queries in queries]
     loop = asyncio.get_event_loop()
@@ -116,10 +126,7 @@ async def bm25_pure(queries: List[str], top_k: int, tokenizer, bm25_api: BM25Oka
     :return: The tuple contains a list of passage ids that retrieved from bm25 and its scores.
     """
     # I don't make queries operation to async, because queries length might be small, so it will occur overhead.
-    if isinstance(tokenizer, PreTrainedTokenizerBase):
-        tokenized_queries = tokenizer(queries).input_ids
-    else:
-        tokenized_queries = tokenizer(queries)
+    tokenized_queries = tokenize(queries, tokenizer)
     id_result = []
     score_result = []
     for query in tokenized_queries:
@@ -139,6 +146,28 @@ async def bm25_pure(queries: List[str], top_k: int, tokenizer, bm25_api: BM25Oka
     return list(id_result), list(score_result)
 
 
+def get_bm25_scores(queries: List[str], ids: List[str], tokenizer,
+                    bm25_api: BM25Okapi, bm25_corpus: Dict) -> List[float]:
+    if len(ids) == 0 or not bool(ids):
+        return []
+    tokenized_queries = tokenize(queries, tokenizer)
+    result_dict = {id_: [] for id_ in ids}
+    for query in tokenized_queries:
+        scores = bm25_api.get_scores(query)
+        for i, id_ in enumerate(ids):
+            result_dict[id_].append(scores[bm25_corpus['passage_id'].index(id_)])
+    result_df = pd.DataFrame(result_dict)
+    return result_df.max(axis=0).tolist()
+
+
+def tokenize(queries: List[str], tokenizer) -> List[List[int]]:
+    if isinstance(tokenizer, PreTrainedTokenizerBase):
+        tokenized_queries = tokenizer(queries).input_ids
+    else:
+        tokenized_queries = tokenizer(queries)
+    return tokenized_queries
+
+
 def bm25_ingest(corpus_path: str, corpus_data: pd.DataFrame, bm25_tokenizer: str = 'porter_stemmer'):
     if not corpus_path.endswith('.pkl'):
         raise ValueError(f"Corpus path {corpus_path} is not a pickle file.")

autorag-0.2.11/autorag/nodes/retrieval/hybrid_cc.py

@@ -0,0 +1,137 @@
+from typing import Tuple, List
+
+import numpy as np
+import pandas as pd
+
+from autorag.nodes.retrieval import retrieval_node
+
+
+def normalize_mm(scores: List[str], fixed_min_value: float = 0):
+    arr = np.array(scores)
+    max_value = np.max(arr)
+    min_value = np.min(arr)
+    norm_score = (arr - min_value) / (max_value - min_value)
+    return norm_score
+
+
+def normalize_tmm(scores: List[str], fixed_min_value: float):
+    arr = np.array(scores)
+    max_value = np.max(arr)
+    norm_score = (arr - fixed_min_value) / (max_value - fixed_min_value)
+    return norm_score
+
+
+def normalize_z(scores: List[str], fixed_min_value: float = 0):
+    arr = np.array(scores)
+    mean_value = np.mean(arr)
+    std_value = np.std(arr)
+    norm_score = (arr - mean_value) / std_value
+    return norm_score
+
+
+def normalize_dbsf(scores: List[str], fixed_min_value: float = 0):
+    arr = np.array(scores)
+    mean_value = np.mean(arr)
+    std_value = np.std(arr)
+    min_value = mean_value - 3 * std_value
+    max_value = mean_value + 3 * std_value
+    norm_score = (arr - min_value) / (max_value - min_value)
+    return norm_score
+
+
+normalize_method_dict = {
+    'mm': normalize_mm,
+    'tmm': normalize_tmm,
+    'z': normalize_z,
+    'dbsf': normalize_dbsf,
+}
+
+
+@retrieval_node
+def hybrid_cc(
+        ids: Tuple,
+        scores: Tuple,
+        top_k: int,
+        weight: float,
+        normalize_method: str = 'mm',
+        semantic_theoretical_min_value: float = -1.0,
+        lexical_theoretical_min_value: float = 0.0,
+) -> Tuple[List[List[str]], List[List[float]]]:
+    """
+    Hybrid CC function.
+    CC (convex combination) is a method to fuse lexical and semantic retrieval results.
+    It is a method that first normalizes the scores of each retrieval result,
+    and then combines them with the given weights.
+    It is uniquer than other retrieval modules, because it does not really execute retrieval,
+    but just fuse the results of other retrieval functions.
+    So you have to run more than two retrieval modules before running this function.
+    And collect ids and scores result from each retrieval module.
+    Make it as tuple and input it to this function.
+
+    :param ids: The tuple of ids that you want to fuse.
+        The length of this must be the same as the length of scores.
+        The semantic retrieval ids must be the first index.
+    :param scores: The retrieve scores that you want to fuse.
+        The length of this must be the same as the length of ids.
+        The semantic retrieval scores must be the first index.
+    :param top_k: The number of passages to be retrieved.
+    :param normalize_method: The normalization method to use.
+      There are some normalization method that you can use at the hybrid cc method.
+      AutoRAG support following.
+        - `mm`: Min-max scaling
+        - `tmm`: Theoretical min-max scaling
+        - `z`: z-score normalization
+        - `dbsf`: 3-sigma normalization
+    :param weight: The weight value. If the weight is 1.0, it means the
+      weight to the semantic module will be 1.0 and weight to the lexical module will be 0.0.
+    :param semantic_theoretical_min_value: This value used by `tmm` normalization method. You can set the
+        theoretical minimum value by yourself. Default is -1.
+    :param lexical_theoretical_min_value: This value used by `tmm` normalization method. You can set the
+        theoretical minimum value by yourself. Default is 0.
+    :return: The tuple of ids and fused scores that fused by CC. Plus, the third element is selected weight value.
+    """
+    assert len(ids) == len(scores), "The length of ids and scores must be the same."
+    assert len(ids) > 1, "You must input more than one retrieval results."
+    assert top_k > 0, "top_k must be greater than 0."
+    assert weight >= 0, "The weight must be greater than 0."
+    assert weight <= 1, "The weight must be less than 1."
+
+    df = pd.DataFrame({
+        'semantic_ids': ids[0],
+        'lexical_ids': ids[1],
+        'semantic_score': scores[0],
+        'lexical_score': scores[1],
+    })
+
+    def cc_pure_apply(row):
+        return fuse_per_query(row['semantic_ids'], row['lexical_ids'],
+                              row['semantic_score'], row['lexical_score'],
+                              normalize_method=normalize_method,
+                              weight=weight, top_k=top_k,
+                              semantic_theoretical_min_value=semantic_theoretical_min_value,
+                              lexical_theoretical_min_value=lexical_theoretical_min_value)
+
+    # fixed weight
+    df[['cc_id', 'cc_score']] = df.apply(lambda row: cc_pure_apply(row), axis=1,
+                                         result_type='expand')
+    return df['cc_id'].tolist(), df['cc_score'].tolist()
+
+
+def fuse_per_query(semantic_ids: List[str], lexical_ids: List[str],
+                   semantic_scores: List[float], lexical_scores: List[float],
+                   normalize_method: str,
+                   weight: float,
+                   top_k: int,
+                   semantic_theoretical_min_value: float,
+                   lexical_theoretical_min_value: float):
+    normalize_func = normalize_method_dict[normalize_method]
+    norm_semantic_scores = normalize_func(semantic_scores, semantic_theoretical_min_value)
+    norm_lexical_scores = normalize_func(lexical_scores, lexical_theoretical_min_value)
+    ids = [semantic_ids, lexical_ids]
+    scores = [norm_semantic_scores, norm_lexical_scores]
+    df = pd.concat([pd.Series(dict(zip(_id, score))) for _id, score in zip(ids, scores)], axis=1)
+    df.columns = ['semantic', 'lexical']
+    df = df.fillna(0)
+    df['weighted_sum'] = df.mul((weight, 1.0 - weight)).sum(axis=1)
+    df = df.sort_values(by='weighted_sum', ascending=False)
+    return df.index.tolist()[:top_k], df['weighted_sum'][:top_k].tolist()

{autorag-0.2.9 → autorag-0.2.11}/autorag/nodes/retrieval/hybrid_rrf.py

@@ -10,7 +10,8 @@ def hybrid_rrf(
         ids: Tuple,
         scores: Tuple,
         top_k: int,
-        rrf_k: int = 60) -> Tuple[List[List[str]], List[List[float]]]:
+        weight: int = 60,
+        rrf_k: int = -1, ) -> Tuple[List[List[str]], List[List[float]]]:
     """
     Hybrid RRF function.
     RRF (Rank Reciprocal Fusion) is a method to fuse multiple retrieval results.
@@ -27,15 +28,23 @@ def hybrid_rrf(
     :param scores: The retrieve scores that you want to fuse.
         The length of this must be the same as the length of ids.
     :param top_k: The number of passages to be retrieved.
-    :param rrf_k: Hyperparameter for RRF.
+    :param weight: Hyperparameter for RRF.
+        It was originally rrf_k value.
         Default is 60.
         For more information, please visit our documentation.
+    :param rrf_k: (Deprecated) Hyperparameter for RRF.
+        It was originally rrf_k value. Will remove at further version.
     :return: The tuple of ids and fused scores that fused by RRF.
     """
     assert len(ids) == len(scores), "The length of ids and scores must be the same."
     assert len(ids) > 1, "You must input more than one retrieval results."
     assert top_k > 0, "top_k must be greater than 0."
-    assert rrf_k > 0, "rrf_k must be greater than 0."
+    assert weight > 0, "rrf_k must be greater than 0."
+
+    if rrf_k != -1:
+        weight = int(rrf_k)
+    else:
+        weight = int(weight)
 
     id_df = pd.DataFrame({f'id_{i}': id_list for i, id_list in enumerate(ids)})
     score_df = pd.DataFrame({f'score_{i}': score_list for i, score_list in enumerate(scores)})
@@ -44,9 +53,9 @@ def hybrid_rrf(
     def rrf_pure_apply(row):
         ids_tuple = tuple(row[[f'id_{i}' for i in range(len(ids))]].values)
         scores_tuple = tuple(row[[f'score_{i}' for i in range(len(scores))]].values)
-        return pd.Series(rrf_pure(ids_tuple, scores_tuple, rrf_k, top_k))
+        return pd.Series(rrf_pure(ids_tuple, scores_tuple, weight, top_k))
 
-    df[['rrf_id', 'rrf_score']] = df.swifter.apply(rrf_pure_apply, axis=1)
+    df[['rrf_id', 'rrf_score']] = df.apply(rrf_pure_apply, axis=1)
     return df['rrf_id'].tolist(), df['rrf_score'].tolist()