autorag.nodes.generator package

Submodules

autorag.nodes.generator.base module

autorag.nodes.generator.base.generator_node(func)[source]

autorag.nodes.generator.llama_index_llm module

autorag.nodes.generator.llama_index_llm.llama_index_llm(prompts: List[str], llm: LLM, batch: int = 16) Tuple[List[str], List[List[int]], List[List[float]]][source]

Llama Index LLM module. It gets the LLM instance from llama index, and returns generated text by the input prompt. It does not generate the right log probs, but it returns the pseudo log probs, which is not meant to be used for other modules.

Parameters:
  • prompts – A list of prompts.

  • llm – A llama index LLM instance.

  • batch – The batch size for llm. Set low if you face some errors.

Returns:

A tuple of three elements. The first element is a list of generated text. The second element is a list of generated text’s token ids, used tokenizer is GPT2Tokenizer. The third element is a list of generated text’s pseudo log probs.

autorag.nodes.generator.openai_llm module

async autorag.nodes.generator.openai_llm.get_result(prompt: str, client: AsyncOpenAI, model: str, tokenizer: Encoding, **kwargs)[source]
async autorag.nodes.generator.openai_llm.get_result_o1(prompt: str, client: AsyncOpenAI, model: str, tokenizer: Encoding, **kwargs)[source]
autorag.nodes.generator.openai_llm.openai_llm(prompts: List[str], llm: str = 'gpt-3.5-turbo', batch: int = 16, truncate: bool = True, api_key: str = None, **kwargs) Tuple[List[str], List[List[int]], List[List[float]]][source]

OpenAI generator module. Uses official openai library for generating answer from the given prompt. It returns real token ids and log probs, so you must use this for using token ids and log probs.

Parameters:
  • prompts – A list of prompts.

  • llm – A model name for openai. Default is gpt-3.5-turbo.

  • batch – Batch size for openai api call. If you get API limit errors, you should lower the batch size. Default is 16.

  • truncate – Whether to truncate the input prompt. Default is True.

  • api_key – OpenAI API key. You can set this by passing env variable OPENAI_API_KEY

  • kwargs – The optional parameter for openai api call openai.chat.completion See https://platform.openai.com/docs/api-reference/chat/create for more details.

Returns:

A tuple of three elements. The first element is a list of generated text. The second element is a list of generated text’s token ids. The third element is a list of generated text’s log probs.

autorag.nodes.generator.openai_llm.truncate_by_token(prompt: str, tokenizer: Encoding, max_token_size: int)[source]

autorag.nodes.generator.run module

autorag.nodes.generator.run.evaluate_generator_node(result_df: DataFrame, metric_inputs: List[MetricInput], metrics: List[str] | List[Dict])[source]
autorag.nodes.generator.run.run_generator_node(modules: List[Callable], module_params: List[Dict], previous_result: DataFrame, node_line_dir: str, strategies: Dict) DataFrame[source]

Run evaluation and select the best module among generator node results. And save the results and summary to generator node directory.

Parameters:
  • modules – Generator modules to run.

  • module_params – Generator module parameters. Including node parameters, which is used for every module in this node.

  • previous_result – Previous result dataframe. Could be prompt maker node’s result.

  • node_line_dir – This node line’s directory.

  • strategies – Strategies for generator node.

Returns:

The best result dataframe. It contains previous result columns and generator node’s result columns.

autorag.nodes.generator.vllm module

autorag.nodes.generator.vllm.destroy_vllm_instance(vllm_instance)[source]
autorag.nodes.generator.vllm.make_vllm_instance(llm: str, input_args: Dict)[source]
autorag.nodes.generator.vllm.vllm(prompts: List[str], llm: str, **kwargs) Tuple[List[str], List[List[int]], List[List[float]]][source]

Vllm module. It gets the VLLM instance, and returns generated texts by the input prompt. You can set logprobs to get the log probs of the generated text. Default logprobs is 1.

Parameters:
  • prompts – A list of prompts.

  • llm – Model name of vLLM.

  • kwargs – The extra parameters for generating the text.

Returns:

A tuple of three elements. The first element is a list of generated text. The second element is a list of generated text’s token ids. The third element is a list of generated text’s log probs.

Module contents