trinity.common.models.model module

Base Model Class

class trinity.common.models.model.InferenceModel

Bases: ABC

A model for high performance for rollout inference.

async generate(prompt: str, **kwargs) → Sequence[Experience]

Generate a responses from a prompt in async.

async chat(messages: List[dict], **kwargs) → Sequence[Experience]

Generate experiences from a list of history chat messages in async.

async logprobs(tokens: List[int]) → Tensor

Generate logprobs for a list of tokens in async.

async convert_messages_to_experience(messages: List[dict]) → Experience

Convert a list of messages into an experience in async.

abstract get_model_version() → int

Get the checkpoint version.

get_available_address() → Tuple[str, int]

Get the address of the actor.

class trinity.common.models.model.ModelWrapper(model: Any, model_type: str = 'vllm', enable_history: bool = False)

Bases: object

A wrapper for the InferenceModel Ray Actor

__init__(model: Any, model_type: str = 'vllm', enable_history: bool = False)

generate(*args, **kwargs)

async generate_async(*args, **kwargs)

chat(*args, **kwargs)

async chat_async(*args, **kwargs)

logprobs(tokens: List[int]) → Tensor

Calculate the logprobs of the given tokens.

async logprobs_async(tokens: List[int]) → Tensor

Calculate the logprobs of the given tokens in async.

convert_messages_to_experience(messages: List[dict]) → Experience

Convert a list of messages into an experience.

async convert_messages_to_experience_async(messages: List[dict]) → Experience

Convert a list of messages into an experience in async.

property model_version: int

Get the version of the model.

get_openai_client() → OpenAI

Get the openai client.

Returns:

The openai client. And model_path is added to the client which refers to the model path.

Return type:

openai.OpenAI

extract_experience_from_history(clear_history: bool = True) → List[Experience]

Extract experiences from the history.

trinity.common.models.model.convert_api_output_to_experience(output) → List[Experience]

Convert the API output to a list of experiences.

trinity.common.models.model.extract_logprobs(choice) → Tensor

Extract logprobs from a list of logprob dictionaries.