AutoRAG 0.2.10__tar.gz → 0.2.12__tar.gz

{autorag-0.2.10 → autorag-0.2.12}/AutoRAG.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.10
+Version: 0.2.12
 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
@@ -242,7 +242,7 @@ Requires-Dist: uvicorn
 Requires-Dist: torch
 Requires-Dist: sentencepiece
 Requires-Dist: guidance
-Requires-Dist: cohere>=5.0.0a9
+Requires-Dist: cohere<=5.6.2
 Requires-Dist: tokenlog>=0.0.2
 Requires-Dist: aiohttp
 Requires-Dist: bert_score
@@ -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.10 → autorag-0.2.12}/AutoRAG.egg-info/requires.txt

@@ -20,7 +20,7 @@ uvicorn
 torch
 sentencepiece
 guidance
-cohere>=5.0.0a9
+cohere<=5.6.2
 tokenlog>=0.0.2
 aiohttp
 bert_score
@@ -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.10 → autorag-0.2.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: AutoRAG
-Version: 0.2.10
+Version: 0.2.12
 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
@@ -242,7 +242,7 @@ Requires-Dist: uvicorn
 Requires-Dist: torch
 Requires-Dist: sentencepiece
 Requires-Dist: guidance
-Requires-Dist: cohere>=5.0.0a9
+Requires-Dist: cohere<=5.6.2
 Requires-Dist: tokenlog>=0.0.2
 Requires-Dist: aiohttp
 Requires-Dist: bert_score
@@ -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.12/autorag/VERSION

@@ -0,0 +1 @@
+0.2.12

{autorag-0.2.10 → autorag-0.2.12}/autorag/evaluation/metric/generation.py

@@ -6,13 +6,13 @@ from typing import List, Optional
 
 import evaluate
 import pandas as pd
-from sacrebleu.metrics.bleu import BLEU
 import torch
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
 from openai import AsyncOpenAI
 from rouge_score import tokenizers
 from rouge_score.rouge_scorer import RougeScorer
+from sacrebleu.metrics.bleu import BLEU
 
 from autorag import embedding_models
 from autorag.evaluation.metric.util import calculate_cosine_similarity
@@ -65,7 +65,9 @@ def huggingface_evaluate(instance, key: str,
 
 
 @convert_inputs_to_list
-def bleu(generation_gt: List[List[str]], generations: [str], tokenize: str|None = None, smooth_method: str = 'exp', smooth_value: Optional[float] = None, max_ngram_order: int = 4, trg_lang: str = '', **kwargs) -> List[float]:
+def bleu(generation_gt: List[List[str]], generations: [str], tokenize: Optional[str] = None,
+         smooth_method: str = 'exp', smooth_value: Optional[float] = None, max_ngram_order: int = 4,
+         trg_lang: str = '', **kwargs) -> List[float]:
     """
     Computes the BLEU metric given pred and ground-truth.
 
@@ -79,9 +81,10 @@ def bleu(generation_gt: List[List[str]], generations: [str], tokenize: str|None
         Because it can be a multiple ground truth.
     :param generations: A list of generations that LLM generated.
     """
-    bleu = BLEU(tokenize=tokenize, smooth_method=smooth_method, smooth_value=smooth_value, max_ngram_order=max_ngram_order, trg_lang=trg_lang, **kwargs)
+    bleu_instance = BLEU(tokenize=tokenize, smooth_method=smooth_method, smooth_value=smooth_value,
+                         max_ngram_order=max_ngram_order, trg_lang=trg_lang, **kwargs)
 
-    result = list(map(lambda x: bleu.sentence_score(x[0], x[1]).score, zip(generations, generation_gt)))
+    result = list(map(lambda x: bleu_instance.sentence_score(x[0], x[1]).score, zip(generations, generation_gt)))
     return result
 
 
@@ -319,4 +322,8 @@ def bert_score(generation_gt: List[List[str]], generations: List[str],
     df['bert_score'] = evaluator.compute(predictions=df['prediction'].tolist(),
                                          references=df['reference'].tolist(),
                                          lang=lang, nthreads=n_threads, batch_size=batch)['f1']
+
+    del evaluator
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     return df.groupby(level=0)['bert_score'].max().tolist()

{autorag-0.2.10 → autorag-0.2.12}/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.10 → autorag-0.2.12}/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.10 → autorag-0.2.12}/autorag/nodes/retrieval/run.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import pathlib
+from copy import deepcopy
 from typing import List, Callable, Dict, Tuple
 
 import numpy as np
@@ -9,6 +10,8 @@ from tqdm import tqdm
 
 from autorag.evaluation import evaluate_retrieval
 from autorag.strategy import measure_speed, filter_by_threshold, select_best
+from autorag.support import get_support_modules
+from autorag.utils.util import get_best_row, to_list
 
 logger = logging.getLogger("AutoRAG")
 
@@ -152,7 +155,8 @@ def run_retrieval_node(modules: List[Callable],
             filename_first += len(hybrid_modules)
         else:  # for Evaluator
             # get id and score
-            ids_scores = get_ids_and_scores(save_dir, [semantic_selected_filename, lexical_selected_filename])
+            ids_scores = get_ids_and_scores(save_dir, [semantic_selected_filename, lexical_selected_filename],
+                                            semantic_summary_df, lexical_summary_df, previous_result)
             hybrid_module_params = list(map(lambda x: {**x, **ids_scores}, hybrid_module_params))
 
             # optimize each modules
@@ -237,17 +241,66 @@ def edit_summary_df_params(summary_df: pd.DataFrame, target_modules, target_modu
     return summary_df
 
 
-def get_ids_and_scores(node_dir: str, filenames: List[str]) -> Dict:
+def get_ids_and_scores(node_dir: str, filenames: List[str],
+                       semantic_summary_df: pd.DataFrame,
+                       lexical_summary_df: pd.DataFrame,
+                       previous_result) -> Dict[str, Tuple[List[List[str]], List[List[float]]]]:
+    project_dir = pathlib.PurePath(node_dir).parent.parent.parent
     best_results_df = list(
         map(lambda filename: pd.read_parquet(os.path.join(node_dir, filename), engine='pyarrow'), filenames))
     ids = tuple(map(lambda df: df['retrieved_ids'].apply(list).tolist(), best_results_df))
     scores = tuple(map(lambda df: df['retrieve_scores'].apply(list).tolist(), best_results_df))
+    # search non-duplicate ids
+    semantic_ids = deepcopy(ids[0])
+    lexical_ids = deepcopy(ids[1])
+
+    def get_non_duplicate_ids(target_ids, compare_ids) -> List[List[str]]:
+        """
+        Get non-duplicate ids from target_ids and compare_ids.
+        If you want to non-duplicate ids of semantic_ids, you have to put it at target_ids.
+        """
+        result_ids = []
+        assert len(target_ids) == len(compare_ids)
+        for target_id_list, compare_id_list in zip(target_ids, compare_ids):
+            query_duplicated = list(set(compare_id_list) - set(target_id_list))
+            duplicate_list = query_duplicated if len(query_duplicated) != 0 else []
+            result_ids.append(duplicate_list)
+        return result_ids
+
+    lexical_target_ids = get_non_duplicate_ids(lexical_ids, semantic_ids)
+    semantic_target_ids = get_non_duplicate_ids(semantic_ids, lexical_ids)
+
+    new_id_tuple = ([a + b for a, b in zip(semantic_ids, semantic_target_ids)],
+                    [a + b for a, b in zip(lexical_ids, lexical_target_ids)])
+
+    # search non-duplicate ids' scores
+    new_semantic_scores = get_scores_by_ids(semantic_target_ids, semantic_summary_df, project_dir, previous_result)
+    new_lexical_scores = get_scores_by_ids(lexical_target_ids, lexical_summary_df, project_dir, previous_result)
+
+    new_score_tuple = (
+        [a + b for a, b in zip(scores[0], new_semantic_scores)],
+        [a + b for a, b in zip(scores[1], new_lexical_scores)],
+    )
     return {
-        'ids': ids,
-        'scores': scores,
+        'ids': new_id_tuple,
+        'scores': new_score_tuple,
     }
 
 
+def get_scores_by_ids(ids: List[List[str]],
+                      module_summary_df: pd.DataFrame,
+                      project_dir, previous_result) -> List[List[float]]:
+    module_name = get_best_row(module_summary_df)['module_name']
+    module_params = get_best_row(module_summary_df)['module_params']
+    module = get_support_modules(module_name)
+    result_df = module(project_dir=project_dir, previous_result=previous_result, ids=ids, **module_params)
+    return to_list(result_df['retrieve_scores'].tolist())
+
+
+def find_unique_elems(list1: List[str], list2: List[str]) -> List[str]:
+    return list(set(list1).symmetric_difference(set(list2)))
+
+
 def get_hybrid_execution_times(lexical_summary, semantic_summary) -> float:
     lexical_execution_time = lexical_summary.loc[lexical_summary['is_best'] == True].iloc[0]['execution_time']
     semantic_execution_time = semantic_summary.loc[semantic_summary['is_best'] == True].iloc[0]['execution_time']

{autorag-0.2.10 → autorag-0.2.12}/autorag/nodes/retrieval/vectordb.py

@@ -1,21 +1,23 @@
 import asyncio
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 
 import chromadb
 import pandas as pd
+from chromadb import GetResult, QueryResult
 from chromadb.utils.batch_utils import create_batches
 from llama_index.core.embeddings import BaseEmbedding
 from llama_index.embeddings.openai import OpenAIEmbedding
 
 from autorag.nodes.retrieval.base import retrieval_node, evenly_distribute_passages
 from autorag.utils import validate_corpus_dataset
-from autorag.utils.util import process_batch, openai_truncate_by_token
+from autorag.utils.util import process_batch, openai_truncate_by_token, flatten_apply
 
 
 @retrieval_node
 def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collection,
              embedding_model: BaseEmbedding,
-             embedding_batch: int = 128) -> Tuple[List[List[str]], List[List[float]]]:
+             embedding_batch: int = 128,
+             ids: Optional[List[List[str]]] = None) -> Tuple[List[List[str]], List[List[float]]]:
     """
     VectorDB retrieval function.
     You have to get a chroma collection that is already ingested.
@@ -29,6 +31,9 @@ def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collecti
     :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.
+    :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 vectordb and 2-d list of its scores.
         It will be a length of queries. And each element has a length of top_k.
@@ -36,8 +41,25 @@ def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collecti
     # check if bm25_corpus is valid
     assert (collection.count() > 0), \
         "collection must contain at least one document. Please check you ingested collection correctly."
+    # truncate queries and embedding execution here.
+    openai_embedding_limit = 8191
+    if isinstance(embedding_model, OpenAIEmbedding):
+        queries = list(map(lambda query_list: openai_truncate_by_token(query_list, openai_embedding_limit,
+                                                                       embedding_model.model_name), queries))
+
+    query_embeddings = flatten_apply(run_query_embedding_batch, queries, embedding_model=embedding_model,
+                                     batch_size=embedding_batch)
+
+    # if ids are specified, fetch the ids score from Chroma
+    if ids is not None:
+        client = chromadb.Client()
+        score_result = list(map(
+            lambda query_embedding_list, id_list: get_id_scores(id_list, query_embedding_list, collection, client),
+            query_embeddings, ids))
+        return ids, score_result
+
     # run async vector_db_pure function
-    tasks = [vectordb_pure(input_queries, top_k, collection, embedding_model) for input_queries in queries]
+    tasks = [vectordb_pure(query_embedding, top_k, collection) for query_embedding in query_embeddings]
     loop = asyncio.get_event_loop()
     results = loop.run_until_complete(process_batch(tasks, batch_size=embedding_batch))
     id_result = list(map(lambda x: x[0], results))
@@ -45,30 +67,20 @@ def vectordb(queries: List[List[str]], top_k: int, collection: chromadb.Collecti
     return id_result, score_result
 
 
-async def vectordb_pure(queries: List[str], top_k: int, collection: chromadb.Collection,
-                        embedding_model: BaseEmbedding) -> Tuple[List[str], List[float]]:
+async def vectordb_pure(query_embeddings: List[List[float]], top_k: int, collection: chromadb.Collection) -> Tuple[
+    List[str], List[float]]:
     """
     Async VectorDB retrieval function.
     Its usage is for async retrieval of vector_db row by row.
 
-    :param queries: A list of query strings.
+    :param query_embeddings: A list of query embeddings.
     :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.
 
     :return: The tuple contains a list of passage ids that retrieved from vectordb and a list of its scores.
     """
-    # truncate queries if embedding is openai and the query exceeded token limit
-    openai_embedding_limit = 8191
-    if isinstance(embedding_model, OpenAIEmbedding):
-        queries = openai_truncate_by_token(queries, openai_embedding_limit,
-                                           embedding_model.model_name)
-
-    # embed query
-    embedded_queries = list(map(embedding_model.get_query_embedding, queries))
-
     id_result, score_result = [], []
-    for embedded_query in embedded_queries:
+    for embedded_query in query_embeddings:
         result = collection.query(query_embeddings=embedded_query, n_results=top_k)
         id_result.extend(result['ids'])
         score_result.extend(result['distances'])
@@ -107,3 +119,34 @@ def vectordb_ingest(collection: chromadb.Collection, corpus_data: pd.DataFrame,
             ids = batch[0]
             embed_content = batch[1]
             collection.add(ids=ids, embeddings=embed_content)
+
+
+def run_query_embedding_batch(queries: List[str], embedding_model: BaseEmbedding,
+                              batch_size: int) -> List[List[float]]:
+    result = []
+    for i in range(0, len(queries), batch_size):
+        batch = queries[i:i + batch_size]
+        embeddings = embedding_model.get_text_embedding_batch(batch)
+        result.extend(embeddings)
+    return result
+
+
+def get_id_scores(ids: List[str], query_embeddings: List[List[float]],
+                  collection: chromadb.Collection, temp_client: chromadb.Client) -> List[float]:
+    if len(ids) == 0 or ids is None or not bool(ids):
+        return []
+
+    id_results: GetResult = collection.get(ids, include=['embeddings'])
+    temp_collection = temp_client.create_collection(name="temp", metadata={"hnsw:space": "cosine"})
+    temp_collection.add(ids=id_results['ids'], embeddings=id_results['embeddings'])
+
+    query_result: QueryResult = temp_collection.query(query_embeddings=query_embeddings, n_results=len(ids))
+    assert len(query_result['ids']) == len(query_result['distances'])
+    id_scores_dict = {id_: [] for id_ in ids}
+    for id_list, score_list in zip(query_result['ids'], query_result['distances']):
+        for id_ in list(id_scores_dict.keys()):
+            id_idx = id_list.index(id_)
+            id_scores_dict[id_].append(score_list[id_idx])
+    id_scores_pd = pd.DataFrame(id_scores_dict)
+    temp_client.delete_collection("temp")
+    return id_scores_pd.max(axis=0).tolist()

{autorag-0.2.10 → autorag-0.2.12}/autorag/utils/util.py

@@ -484,3 +484,19 @@ def convert_inputs_to_list(func):
         return func(*new_args, **new_kwargs)
 
     return wrapper
+
+
+def get_best_row(summary_df: pd.DataFrame,
+                 best_column_name: str = 'is_best') -> pd.Series:
+    """
+    From summary dataframe, find the best result row by 'is_best' column and return it.
+
+    :param summary_df: Summary dataframe created by AutoRAG.
+    :param best_column_name: The column name that indicates the best result.
+        Default is 'is_best'.
+        You don't have to change this unless the column name is different.
+    :return: Best row pandas Series instance.
+    """
+    bests = summary_df.loc[summary_df[best_column_name] == True]
+    assert len(bests) == 1, "There must be only one best result."
+    return bests.iloc[0]

{autorag-0.2.10 → autorag-0.2.12}/requirements.txt

@@ -20,7 +20,7 @@ uvicorn  # for api server
 torch  # for monot5 reranker
 sentencepiece  # for monot5 reranker
 guidance # for qa data creation
-cohere>=5.0.0a9 # for cohere services
+cohere<=5.6.2 # for cohere services
 tokenlog>=0.0.2 # for token logging
 aiohttp # for async http requests
 bert_score # for bert score
@@ -39,7 +39,7 @@ llama-index-readers-file
 llama-index-embeddings-openai
 llama-index-embeddings-huggingface
 # LLMs
-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.10 → autorag-0.2.12}/tests/autorag/evaluate/metric/test_generation_metric.py

@@ -47,6 +47,7 @@ def ko_base_test_generation_metrics(func, solution, **kwargs):
     assert all(list(map(lambda x: x[0] == pytest.approx(x[1], 0.001),
                         zip(scores, solution))))
 
+
 def test_bleu():
     base_test_generation_metrics(bleu, [51.1507, 23.5783, 100.0], lowercase=True)
 

{autorag-0.2.10 → autorag-0.2.12}/tests/autorag/nodes/generator/test_openai.py

@@ -12,7 +12,7 @@ from tests.mock import mock_openai_chat_create
 @patch.object(openai.resources.chat.completions.AsyncCompletions, 'create', mock_openai_chat_create)
 def test_openai_llm():
     openai_original = openai_llm.__wrapped__
-    model = "gpt-3.5-turbo"
+    model = "gpt-4o-mini"
     answers, tokens, log_probs = openai_original(prompts, model, batch=1, temperature=0.5, logprobs=False, n=3)
     check_generated_texts(answers)
     check_generated_tokens(tokens)
@@ -26,7 +26,7 @@ def test_openai_llm_node():
             'prompts': prompts,
             'qid': ['id-1', 'id-2', 'id-3']
         })
-    result_df = openai_llm(project_dir='.', previous_result=previous_result, llm='gpt-4o')
+    result_df = openai_llm(project_dir='.', previous_result=previous_result, llm='gpt-4o-mini')
     check_generated_texts(result_df['generated_texts'].tolist())
     check_generated_tokens(result_df['generated_tokens'].tolist())
     check_generated_log_probs(result_df['generated_log_probs'].tolist())

{autorag-0.2.10 → autorag-0.2.12}/tests/autorag/nodes/retrieval/test_bm25.py

@@ -7,8 +7,10 @@ import pytest
 
 from autorag.nodes.retrieval import bm25
 from autorag.nodes.retrieval.bm25 import bm25_ingest, tokenize_ko_kiwi, tokenize_porter_stemmer, tokenize_space
+from autorag.utils.util import to_list
 from tests.autorag.nodes.retrieval.test_retrieval_base import (queries, project_dir, corpus_df, previous_result,
-                                                               base_retrieval_test, base_retrieval_node_test)
+                                                               base_retrieval_test, base_retrieval_node_test,
+                                                               searchable_input_ids)
 
 
 @pytest.fixture
@@ -27,12 +29,53 @@ def test_bm25_retrieval(ingested_bm25_path):
     base_retrieval_test(id_result, score_result, top_k)
 
 
+def test_bm25_retrieval_ids(ingested_bm25_path):
+    with open(ingested_bm25_path, 'rb') as r:
+        bm25_corpus = pickle.load(r)
+    original_bm25 = bm25.__wrapped__
+    input_ids = [["doc2", "doc3"],
+                 ["doc1"],
+                 ["doc3", "doc4"]]
+    id_result, score_result = original_bm25(queries, top_k=3, bm25_corpus=bm25_corpus,
+                                            ids=input_ids)
+    assert id_result == input_ids
+    assert len(score_result) == 3
+    assert len(score_result[0]) == 2
+    assert len(score_result[1]) == 1
+    assert len(score_result[2]) == 2
+
+
+def test_bm25_retrieval_ids_empty(ingested_bm25_path):
+    with open(ingested_bm25_path, 'rb') as r:
+        bm25_corpus = pickle.load(r)
+    original_bm25 = bm25.__wrapped__
+    input_ids = [["doc2", "doc3"],
+                 [],
+                 ["doc3"]]
+    id_result, score_result = original_bm25(queries, top_k=3, bm25_corpus=bm25_corpus,
+                                            ids=input_ids)
+    assert id_result == input_ids
+    assert len(score_result) == 3
+    assert len(score_result[0]) == 2
+    assert len(score_result[1]) == 0
+    assert len(score_result[2]) == 1
+
+
 def test_bm25_node():
     result_df = bm25(project_dir=project_dir, previous_result=previous_result, top_k=4,
                      bm25_tokenizer='gpt2')
     base_retrieval_node_test(result_df)
 
 
+def test_bm25_node_ids():
+    result_df = bm25(project_dir=project_dir, previous_result=previous_result, top_k=4,
+                     bm25_tokenizer='gpt2', ids=searchable_input_ids)
+    assert to_list(result_df['retrieved_ids'].tolist()) == searchable_input_ids
+    score_result = to_list(result_df['retrieve_scores'].tolist())
+    assert len(score_result) == 5
+    assert len(score_result[0]) == 2
+
+
 def test_bm25_ingest(ingested_bm25_path):
     with open(ingested_bm25_path, 'rb') as r:
         corpus = pickle.load(r)

{autorag-0.2.10 → autorag-0.2.12}/tests/autorag/nodes/retrieval/test_retrieval_base.py

@@ -1,6 +1,5 @@
 import os
 import pathlib
-
 from datetime import datetime
 
 import pandas as pd
@@ -27,6 +26,13 @@ contents = ["This is a test document 1.", "This is a test document 2.", "This is
 metadata = [{'datetime': datetime.now()} for _ in range(5)]
 corpus_df = pd.DataFrame({"doc_id": doc_id, "contents": contents, "metadata": metadata})
 
+searchable_input_ids = [["aba8293d-2eb9-4a35-8bda-e6c87794c28d", "6f20af70-48b7-4171-a8d7-967ea583a595"],
+                        ["5efb3ac8-04f6-4184-94b3-61cbce080e86", "191c54df-703a-477d-86b6-99183f254799"],
+                        ["99f0e6df-3f03-4bdb-ba24-8f387a783c55", "5b957791-3c7b-4f29-a410-8d005a538855"],
+                        ["fbb8e444-b487-4ba7-9d0b-75b243fd666b", "5efb3ac8-04f6-4184-94b3-61cbce080e86"],
+                        ["1950c5e6-ee53-4b0f-9b32-e01d91d6ce9c", "32207872-dee5-4613-887e-b74fb5c36457"],
+                        ]
+
 
 def base_retrieval_test(id_result, score_result, top_k):
     assert len(id_result) == len(score_result) == 3