Repository: Marker-Inc-Korea/AutoRAG

Issue:
Use generator at hyde and decompose module

Discussion:

Patch:
diff --git a/autorag/nodes/queryexpansion/base.py b/autorag/nodes/queryexpansion/base.py
index ac55fb7..05dcbd5 100644
--- a/autorag/nodes/queryexpansion/base.py
+++ b/autorag/nodes/queryexpansion/base.py
@@ -1,11 +1,11 @@
 import functools
 import logging
 from pathlib import Path
-from typing import List, Union
+from typing import List, Union, Dict, Optional
 
 import pandas as pd
 
-from autorag import generator_models
+from autorag.support import get_support_modules
 from autorag.utils import result_to_dataframe, validate_qa_dataset
 
 logger = logging.getLogger("AutoRAG")
@@ -28,31 +28,33 @@ def query_expansion_node(func):
         if func.__name__ == "pass_query_expansion":
             return func(queries=queries)
 
-        # set module parameters
-        llm_str = kwargs.pop("llm")
-
         # pop prompt from kwargs
         if "prompt" in kwargs.keys():
             prompt = kwargs.pop("prompt")
         else:
             prompt = ""
 
-        # pop batch from kwargs
-        if "batch" in kwargs.keys():
-            batch = kwargs.pop("batch")
-        else:
-            batch = 16
-
-        # set llm model for query expansion
-        if llm_str in generator_models:
-            llm = generator_models[llm_str](**kwargs)
-        else:
-            logger.error(f"llm_str {llm_str} does not exist.")
-            raise KeyError(f"llm_str {llm_str} does not exist.")
+        # set generator module for query expansion
+        generator_callable, generator_param = make_generator_callable_param(kwargs)
 
         # run query expansion function
-        expanded_queries = func(queries=queries, llm=llm, prompt=prompt, batch=batch)
-        del llm
+        expanded_queries = func(queries=queries,
+                                prompt=prompt,
+                                generator_func=generator_callable,
+                                generator_params=generator_param)
         return expanded_queries
 
     return wrapper
+
+
+def make_generator_callable_param(generator_dict: Optional[Dict]):
+    if 'generator_module_type' not in generator_dict.keys():
+        generator_dict = {
+            'generator_module_type': 'llama_index_llm',
+            'llm': 'openai',
+            'model': 'gpt-3.5-turbo',
+        }
+    module_str = generator_dict.pop('generator_module_type')
+    module_callable = get_support_modules(module_str)
+    module_param = generator_dict
+    return module_callable, module_param
diff --git a/autorag/nodes/queryexpansion/hyde.py b/autorag/nodes/queryexpansion/hyde.py
index ebeac60..24dc948 100644
--- a/autorag/nodes/queryexpansion/hyde.py
+++ b/autorag/nodes/queryexpansion/hyde.py
@@ -1,40 +1,31 @@
-import asyncio
-from typing import List
+from typing import List, Dict, Callable
 
-from llama_index.core.service_context_elements.llm_predictor import LLMPredictorType
+import pandas as pd
 
 from autorag.nodes.queryexpansion.base import query_expansion_node
-from autorag.utils.util import process_batch
 
 hyde_prompt = "Please write a passage to answer the question"
 
 
 @query_expansion_node
-def hyde(queries: List[str], llm: LLMPredictorType,
-         prompt: str = hyde_prompt,
-         batch: int = 16) -> List[List[str]]:
+def hyde(queries: List[str],
+         generator_func: Callable,
+         generator_params: Dict,
+         prompt: str = hyde_prompt) -> List[List[str]]:
     """
     HyDE, which inspired by "Precise Zero-shot Dense Retrieval without Relevance Labels" (https://arxiv.org/pdf/2212.10496.pdf)
     LLM model creates a hypothetical passage.
     And then, retrieve passages using hypothetical passage as a query.
     :param queries: List[str], queries to retrieve.
-    :param llm: llm to use for hypothetical passage generation.
+    :param generator_func: Callable, generator functions.
+    :param generator_params: Dict, generator parameters.
     :param prompt: prompt to use when generating hypothetical passage
-    :param batch: Batch size for llm.
-        Default is 16.
     :return: List[List[str]], List of hyde results.
     """
-    # Run async query_decompose_pure function
-    tasks = [hyde_pure(query, llm, prompt) for query in queries]
-    loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+    full_prompts = list(
+        map(lambda x: (prompt if not bool(prompt) else hyde_prompt) + f"\nQuestion: {x}\nPassage:", queries))
+    input_df = pd.DataFrame({"prompts": full_prompts})
+    result_df = generator_func(project_dir=None, previous_result=input_df, **generator_params)
+    answers = result_df['generated_texts'].tolist()
+    results = list(map(lambda x: [x], answers))
     return results
-
-
-async def hyde_pure(query: str, llm: LLMPredictorType,
-                    prompt: str = hyde_prompt) -> List[str]:
-    if prompt is "":
-        prompt = hyde_prompt
-    full_prompt = prompt + f"\nQuestion: {query}\nPassage:"
-    hyde_answer = await llm.acomplete(full_prompt)
-    return [hyde_answer.text]
diff --git a/autorag/nodes/queryexpansion/query_decompose.py b/autorag/nodes/queryexpansion/query_decompose.py
index 433c5f6..1bc0df8 100644
--- a/autorag/nodes/queryexpansion/query_decompose.py
+++ b/autorag/nodes/queryexpansion/query_decompose.py
@@ -1,10 +1,8 @@
-import asyncio
-from typing import List
+from typing import List, Callable, Dict
 
-from llama_index.core.service_context_elements.llm_predictor import LLMPredictorType
+import pandas as pd
 
 from autorag.nodes.queryexpansion.base import query_expansion_node
-from autorag.utils.util import process_batch
 
 decompose_prompt = """Decompose a question in self-contained sub-questions. Use \"The question needs no decomposition\" when no decomposition is needed.
 
@@ -54,44 +52,45 @@ decompose_prompt = """Decompose a question in self-contained sub-questions. Use
 
 
 @query_expansion_node
-def query_decompose(queries: List[str], llm: LLMPredictorType,
-                    prompt: str = decompose_prompt,
-                    batch: int = 16) -> List[List[str]]:
+def query_decompose(queries: List[str],
+                    generator_func: Callable,
+                    generator_params: Dict,
+                    prompt: str = decompose_prompt) -> List[List[str]]:
     """
     decompose query to little piece of questions.
     :param queries: List[str], queries to decompose.
-    :param llm: LLMPredictorType, language model to use.
+    :param generator_func: Callable, generator functions.
+    :param generator_params: Dict, generator parameters.
     :param prompt: str, prompt to use for query decomposition.
         default prompt comes from Visconde's StrategyQA few-shot prompt.
-    :param batch: int, batch size for llm.
-        Default is 16.
     :return: List[List[str]], list of decomposed query. Return input query if query is not decomposable.
     """
     # Run async query_decompose_pure function
-    tasks = [query_decompose_pure(query, llm, prompt) for query in queries]
-    loop = asyncio.get_event_loop()
-    results = loop.run_until_complete(process_batch(tasks, batch_size=batch))
+    full_prompts = []
+    for query in queries:
+        if bool(prompt):
+            full_prompt = f"prompt: {prompt}\n\n question: {query}"
+        else:
+            full_prompt = decompose_prompt.format(question=query)
+        full_prompts.append(full_prompt)
+    input_df = pd.DataFrame({"prompts": full_prompts})
+    result_df = generator_func(project_dir=None, previous_result=input_df, **generator_params)
+    answers = result_df['generated_texts'].tolist()
+    results = list(map(lambda x: get_query_decompose(x[0], x[1]), zip(queries, answers)))
     return results
 
 
-async def query_decompose_pure(query: str, llm: LLMPredictorType,
-                               prompt: str = decompose_prompt) -> List[str]:
+def get_query_decompose(query: str, answer: str) -> List[str]:
     """
     decompose query to little piece of questions.
     :param query: str, query to decompose.
-    :param llm: LLMPredictorType, language model to use.
-    :param prompt: str, prompt to use for query decomposition.
-        default prompt comes from Visconde's StrategyQA few-shot prompt.
+    :param answer: str, answer from query_decompose function.
     :return: List[str], list of a decomposed query. Return input query if query is not decomposable.
     """
-    if prompt == "":
-        prompt = decompose_prompt
-    full_prompt = "prompt: " + prompt + "\n\n" "question: " + query
-    answer = await llm.acomplete(full_prompt)
-    if answer.text == "the question needs no decomposition.":
+    if answer.lower() == "the question needs no decomposition":
         return [query]
     try:
-        lines = [line.strip() for line in answer.text.splitlines() if line.strip()]
+        lines = [line.strip() for line in answer.splitlines() if line.strip()]
         if lines[0].startswith("Decompositions:"):
             lines.pop(0)
         questions = [line.split(':', 1)[1].strip() for line in lines if ':' in line]
diff --git a/docs/source/nodes/query_expansion/hyde.md b/docs/source/nodes/query_expansion/hyde.md
index a3de9f0..50eae66 100644
--- a/docs/source/nodes/query_expansion/hyde.md
+++ b/docs/source/nodes/query_expansion/hyde.md
@@ -4,19 +4,24 @@ The HyDE is inspired by the paper "[Precise Zero-shot Dense Retrieval without Re
 
 ## **Module Parameters**
 
-**llm**: The query expansion node requires setting parameters related to the Large Language Model (LLM) being used. This includes specifying the LLM provider (e.g., `openai` or a list of providers like `[openai, huggingfacellm]` and the model configuration. By default, if only `openai` is specified without a model, the system uses the default model set in `llama_index`, which is `gpt-3.5-turbo`.
-```{tip}
-Information about the LLM model can be found [Supporting LLM models](../../local_model.md#supporting-llm-models).
-```
-- **Additional Parameters**: 
-  - **batch**: How many calls to make at once. Default is 16.
-  - Other LLM-related parameters such as `model`, `temperature`, and `max_token` can be set. These are passed as keyword arguments (`kwargs`) to the LLM object, allowing for further customization of the LLM's behavior.
+**llm**: The query expansion node requires setting parameters related to our generator modules.
+
+- **generator_module_type**: The type of the generator module to use.
+- **llm**: The type of llm.
+- Other LLM-related parameters such as `model`, `temperature`, and `max_token` can be set. These are passed as keyword
+  arguments (`kwargs`) to the LLM object, allowing for further customization of the LLM's behavior.
+
+**Additional Parameters**:
 
+- **prompt**: You can use your own custom prompt for the LLM model.
+  Default prompt is come from the
+  paper "[Precise Zero-shot Dense Retrieval without Relevance Labels](https://arxiv.org/abs/2212.10496)".
 
 ## **Example config.yaml**
 ```yaml
 modules:
-- module_type: hyde
-  llm: openai
-  max_token: 64
+  - module_type: hyde
+    generator_module_type: llama_index_llm
+    llm: openai
+    max_token: 64
 ```
diff --git a/docs/source/nodes/query_expansion/query_decompose.md b/docs/source/nodes/query_expansion/query_decompose.md
index ed5ce78..ad16315 100644
--- a/docs/source/nodes/query_expansion/query_decompose.md
+++ b/docs/source/nodes/query_expansion/query_decompose.md
@@ -4,19 +4,23 @@ The `query_decompose` is used to decompose a ‘multi-hop question’ into ‘mu
 
 ## **Module Parameters**
 
-**llm**: The query expansion node requires setting parameters related to the Large Language Model (LLM) being used. This includes specifying the LLM provider (e.g., `openai` or a list of providers like `[openai, huggingfacellm]`) and the model configuration. By default, if only `openai` is specified without a model, the system uses the default model set in `llama_index`, which is `gpt-3.5-turbo`.
-```{tip}
-Information about the LLM model can be found [Supporting LLM models](../../local_model.md#supporting-llm-models).
-```
-- **Additional Parameters**: 
-  - **batch**: How many llm calls to make at once. Default is 16.
-  - Other LLM-related parameters such as `model`, `temperature`, and `max_token` can be set. These are passed as keyword arguments (`kwargs`) to the LLM object, allowing for further customization of the LLM's behavior. You can find these parameters at LlamaIndex docs.
+**llm**: The query expansion node requires setting parameters related to our generator modules.
+
+- **generator_module_type**: The type of the generator module to use.
+- **llm**: The type of llm.
+- Other LLM-related parameters such as `model`, `temperature`, and `max_token` can be set. These are passed as keyword
+  arguments (`kwargs`) to the LLM object, allowing for further customization of the LLM's behavior.
+
+**Additional Parameters**:
 
+- **prompt**: You can use your own custom prompt for the LLM model.
+  default prompt comes from Visconde's StrategyQA few-shot prompt.
 
 ## **Example config.yaml**
 ```yaml
 modules:
 - module_type: query_decompose
+  generator_module_type: llama_index_llm
   llm: openai
-  temperature: [0.2, 1.0]
+  model: [ gpt-3.5-turbo-16k, gpt-3.5-turbo-1106 ]
 ```
diff --git a/sample_config/full.yaml b/sample_config/full.yaml
index 9b117e8..ef0ce41 100644
--- a/sample_config/full.yaml
+++ b/sample_config/full.yaml
@@ -14,10 +14,13 @@ node_lines:
       modules:
         - module_type: pass_query_expansion
         - module_type: query_decompose
+          generator_module_type: llama_index_llm
           llm: openai
-          temperature: [0.2, 1.0]
+          model: [ gpt-3.5-turbo-16k, gpt-3.5-turbo-1106 ]
         - module_type: hyde
+          generator_module_type: llama_index_llm
           llm: openai
+          model: [ gpt-3.5-turbo-16k ]
           max_token: 64
 - node_line_name: retrieve_node_line  # Arbitrary node line name
   nodes:


Test patch:
diff --git a/tests/autorag/nodes/queryexpansion/test_hyde.py b/tests/autorag/nodes/queryexpansion/test_hyde.py
index 8091b81..2a7754d 100644
--- a/tests/autorag/nodes/queryexpansion/test_hyde.py
+++ b/tests/autorag/nodes/queryexpansion/test_hyde.py
@@ -1,22 +1,25 @@
-from autorag import generator_models
 from autorag.nodes.queryexpansion import hyde
+from autorag.support import get_support_modules
 from tests.autorag.nodes.queryexpansion.test_query_expansion_base import project_dir, previous_result, \
     base_query_expansion_node_test, ingested_vectordb_node
-from tests.mock import MockLLM
 
 sample_query = ["How many members are in Newjeans?", "What is visconde structure?"]
 
 
 def test_hyde():
-    llm = MockLLM()
+    generator_func = get_support_modules('llama_index_llm')
+    generator_params = {'llm': 'mock'}
     original_hyde = hyde.__wrapped__
-    result = original_hyde(sample_query, llm, prompt="")
+    result = original_hyde(sample_query, generator_func, generator_params, prompt="")
     assert len(result[0]) == 1
     assert len(result) == 2
 
 
 def test_hyde_node(ingested_vectordb_node):
-    generator_models['mock'] = MockLLM
+    generator_dict = {
+        'generator_module_type': 'llama_index_llm',
+        'llm': 'mock'
+    }
     result_df = hyde(project_dir=project_dir, previous_result=previous_result,
-                     llm="mock", max_tokens=64)
+                     **generator_dict)
     base_query_expansion_node_test(result_df)
diff --git a/tests/autorag/nodes/queryexpansion/test_query_decompose.py b/tests/autorag/nodes/queryexpansion/test_query_decompose.py
index a7e456d..3c5ce0c 100644
--- a/tests/autorag/nodes/queryexpansion/test_query_decompose.py
+++ b/tests/autorag/nodes/queryexpansion/test_query_decompose.py
@@ -1,23 +1,26 @@
-from autorag import generator_models
 from autorag.nodes.queryexpansion import query_decompose
+from autorag.support import get_support_modules
 from tests.autorag.nodes.queryexpansion.test_query_expansion_base import project_dir, previous_result, \
     base_query_expansion_node_test, ingested_vectordb_node
-from tests.mock import MockLLM
 
 sample_query = ["Which group has more members, Newjeans or Aespa?", "Which group has more members, STAYC or Aespa?"]
 
 
 def test_query_decompose():
-    llm = MockLLM(temperature=0.2)
+    generator_func = get_support_modules('llama_index_llm')
+    generator_params = {'llm': 'mock'}
     original_query_decompose = query_decompose.__wrapped__
-    result = original_query_decompose(sample_query, llm, prompt="")
+    result = original_query_decompose(sample_query, generator_func, generator_params, prompt="")
     assert len(result[0]) > 1
     assert len(result) == 2
     assert isinstance(result[0][0], str)
 
 
 def test_query_decompose_node(ingested_vectordb_node):
-    generator_models['mock'] = MockLLM
+    generator_dict = {
+        'generator_module_type': 'llama_index_llm',
+        'llm': 'mock'
+    }
     result_df = query_decompose(project_dir=project_dir, previous_result=previous_result,
-                                llm="mock", temperature=0.2)
+                                **generator_dict)
     base_query_expansion_node_test(result_df)
diff --git a/tests/autorag/nodes/queryexpansion/test_query_expansion_run.py b/tests/autorag/nodes/queryexpansion/test_query_expansion_run.py
index 8338710..445a811 100644
--- a/tests/autorag/nodes/queryexpansion/test_query_expansion_run.py
+++ b/tests/autorag/nodes/queryexpansion/test_query_expansion_run.py
@@ -2,17 +2,18 @@ import os.path
 import pathlib
 import shutil
 import tempfile
+from unittest.mock import patch
 
 import pandas as pd
 import pytest
+from llama_index.core.base.llms.types import CompletionResponse
+from llama_index.llms.openai import OpenAI
 
-from autorag import generator_models
 from autorag.nodes.queryexpansion import query_decompose, hyde
 from autorag.nodes.queryexpansion.run import evaluate_one_query_expansion_node
 from autorag.nodes.queryexpansion.run import run_query_expansion_node
 from autorag.nodes.retrieval import bm25
 from autorag.utils.util import load_summary_file
-from tests.mock import MockLLM
 
 root_dir = pathlib.PurePath(os.path.dirname(os.path.realpath(__file__))).parent.parent.parent
 resources_dir = os.path.join(root_dir, "resources")
@@ -76,7 +77,7 @@ def base_query_expansion_test(best_result, node_line_dir):
     assert len(summary_df) == 2
     assert summary_df['filename'][0] == "0.parquet"
     assert summary_df['module_name'][0] == "query_decompose"
-    assert summary_df['module_params'][0] == {'llm': "mock", 'temperature': 0.2, 'batch': 7}
+    assert summary_df['module_params'][0] == {'generator_module_type': 'llama_index_llm', 'llm': "mock", 'batch': 7}
     assert summary_df['execution_time'][0] > 0
     assert summary_df['is_best'][0] == True or summary_df['is_best'][0] == False  # is_best is np.bool_
     # test the best file is saved properly
@@ -88,13 +89,13 @@ def base_query_expansion_test(best_result, node_line_dir):
 
 
 def test_run_query_expansion_node(node_line_dir):
-    generator_models['mock'] = MockLLM
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
     qa_path = os.path.join(project_dir, "data", "qa.parquet")
     previous_result = pd.read_parquet(qa_path)
 
     modules = [query_decompose, hyde]
-    module_params = [{'llm': "mock", 'temperature': 0.2, 'batch': 7}, {'llm': "mock"}]
+    module_params = [{'generator_module_type': 'llama_index_llm', 'llm': "mock", 'batch': 7},
+                     {'generator_module_type': 'llama_index_llm', 'llm': "mock"}]
     strategies = {
         'metrics': metrics,
         'speed_threshold': 5,
@@ -106,13 +107,13 @@ def test_run_query_expansion_node(node_line_dir):
 
 
 def test_run_query_expansion_node_default(node_line_dir):
-    generator_models['mock'] = MockLLM
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
     qa_path = os.path.join(project_dir, "data", "qa.parquet")
     previous_result = pd.read_parquet(qa_path)
 
     modules = [query_decompose, hyde]
-    module_params = [{'llm': "mock", 'temperature': 0.2, 'batch': 7}, {'llm': "mock"}]
+    module_params = [{'generator_module_type': 'llama_index_llm', 'llm': "mock", 'batch': 7},
+                     {'generator_module_type': 'llama_index_llm', 'llm': "mock"}]
     strategies = {
         'metrics': metrics
     }
@@ -121,13 +122,41 @@ def test_run_query_expansion_node_default(node_line_dir):
 
 
 def test_run_query_expansion_one_module(node_line_dir):
-    generator_models['mock'] = MockLLM
     project_dir = pathlib.PurePath(node_line_dir).parent.parent
     qa_path = os.path.join(project_dir, "data", "qa.parquet")
     previous_result = pd.read_parquet(qa_path)
 
     modules = [query_decompose]
-    module_params = [{'llm': "mock", 'temperature': 0.2}]
+    module_params = [{'generator_module_type': 'llama_index_llm', 'llm': "mock"}]
+    strategies = {
+        'metrics': metrics
+    }
+    best_result = run_query_expansion_node(modules, module_params, previous_result, node_line_dir, strategies)
+    assert set(best_result.columns) == {
+        'qid', 'query', 'generation_gt', 'retrieval_gt', 'queries'  # automatically skip evaluation
+    }
+    summary_filepath = os.path.join(node_line_dir, "query_expansion", "summary.csv")
+    assert os.path.exists(summary_filepath)
+    summary_df = load_summary_file(summary_filepath)
+    assert set(summary_df) == {
+        'filename', 'module_name', 'module_params', 'execution_time', 'is_best'
+    }
+    best_filepath = os.path.join(node_line_dir, "query_expansion", f"best_{summary_df['filename'].values[0]}")
+    assert os.path.exists(best_filepath)
+
+
+async def mock_acomplete(self, messages, **kwargs):
+    return CompletionResponse(text=messages)
+
+
+@patch.object(OpenAI, "acomplete", mock_acomplete)
+def test_run_query_expansion_no_generator(node_line_dir):
+    project_dir = pathlib.PurePath(node_line_dir).parent.parent
+    qa_path = os.path.join(project_dir, "data", "qa.parquet")
+    previous_result = pd.read_parquet(qa_path)
+
+    modules = [query_decompose]
+    module_params = [{}]
     strategies = {
         'metrics': metrics
     }
diff --git a/tests/resources/full.yaml b/tests/resources/full.yaml
index 7a62628..0d62078 100644
--- a/tests/resources/full.yaml
+++ b/tests/resources/full.yaml
@@ -11,11 +11,12 @@ node_lines:
           - module_type: vectordb
             embedding_model: openai
       modules:
-        - module_type: pass_query_expansion
         - module_type: query_decompose
+          generator_module_type: llama_index_llm
           llm: openai
-          temperature: [0.2, 1.0]
+          model: [ gpt-3.5-turbo-16k, gpt-3.5-turbo-1106 ]
         - module_type: hyde
+          generator_module_type: llama_index_llm
           llm: openai
           max_token: 64
 - node_line_name: retrieve_node_line  # Arbitrary node line name