`services.nemo_client`#

Module Contents#

A wrapper around NemoQueryLLM for querying models in synthetic data generation

class services.nemo_client.NemoDeployClient(nemo_deploy: NemoQueryLLM)#

A wrapper around NemoQueryLLM for querying models in synthetic data generation

Initialization

query_model( *, messages: collections.abc.Iterable, model: str, conversation_formatter: nemo_curator.services.conversation_formatter.ConversationFormatter | None = None, max_tokens: int | None = None, n: int | None = None, seed: int | None = None, stop: str | None | list[str] = [], stream: bool = False, temperature: float | None = None, top_k: int | None = None, top_p: float | None = None, ) → list[str]#

query_reward_model( *, messages: collections.abc.Iterable, model: str, ) → dict#: Prompts an LLM Reward model to score a conversation between a user and assistant Args: messages: The conversation to calculate a score for. Should be formatted like: [{“role”: “user”, “content”: “Write a sentence”}, {“role”: “assistant”, “content”: “This is a sentence”}, …] model: The name of the model that should be used to calculate the reward. Must be a reward model, cannot be a regular LLM. Returns: A mapping of score_name -> score