Use generator module at hyde and decompose instead of LLMPredictorTyp…

…e from LlamaIndex (#464)

* just commit

* implement generator func at hyde and decompose

* add docs for module parameter

* add no generator test and patch

* change to simple.yaml

* add prompt parameter at docs
delete batch parameter

* use bool and f-string at query_decompose.py

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 wrapper(
         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

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]

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 @@
 
 
 @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]

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
 ```

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 ]
 ```

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:

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)

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)