🦜 Models

class text_machina.src.models.base.TextGenerationModel(model_config)

Bases: ABC

Base class for LLMs.

abstract generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

generate_completions(prompts, generation_config)

Generates a completion for each prompt in a list of prompts.

Parameters:

• prompts (List[str]) – List of prompts to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

List of generated completions.

Return type:

List[str]

class text_machina.src.models.ai21.AI21Model(model_config)

Bases: TextGenerationModel

Generates completions using AI21 models.

Requires the definition of the AI21_API_KEY=<api_key> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.anthropic.AnthropicModel(model_config)

Bases: TextGenerationModel

Generates completions using Anthropic models.

Requires the definition of the ANTRHOPIC_API_KEY=<key> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.azure_openai.AzureOpenAIModel(model_config)

Bases: OpenAIModel

Generates completions using Azure OpenAI models.

Requires the definition of the AZURE_OPENAI_API_KEY=<key> env variable.

class text_machina.src.models.bedrock.BedrockModel(model_config)

Bases: TextGenerationModel

Generates completions using AWS Bedrock models.

Requires the definition of the AWS_ACCESS_KEY_ID=<key> and AWS_SECRET_ACCESS_KEY=<key> environment variables.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

get_completion_from_response_body(response_body)

Obtains the completions from a response body returned by a bedrock model.

Considers the different API schemas that each model provider uses.

Parameters:

response_body (Dict) – the body returned by models in bedrock.

Returns:

the completion of the model extracted from the body.

Return type:

str

get_request_body(prompt, generation_config)

Prepares the request body for a request to a bedrock model.

Considers the different parameters that each model provider accepts.

Parameters:

• prompt (str) – the prompt to use for generating text.

• generationc_config (Dict[str, Any]) – the generation config.

Returns:

a serializable provider-specific request body.

Return type:

Dict

class text_machina.src.models.cohere.CohereModel(model_config)

Bases: TextGenerationModel

Generates completions using Cohere models.

Requires the definition of the COHERE_API_KEY=<key> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.hf_local.HuggingFaceLocalModel(model_config)

Bases: TextGenerationModel

Generates completions using HuggingFace’s models locally deployed.

generate_completion(prompt, generation_config)

Override generate_completions for completeness. This method is not used, since generations are done with batches using generate_completions.

Return type:

str

generate_completions(prompts, generation_config)

Overriden method to generate completions using HuggingFace’s generate method with batches

Return type:

List[str]

class text_machina.src.models.hf_remote.HuggingFaceRemoteModel(model_config)

Bases: TextGenerationModel

Generates completions using HuggingFace’s models remotely deployed (HuggingFace’s Inference API or Inference Endpoints).

Requires the definition of the HF_TOKEN=<token> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.inference_server.InferenceServerModel(model_config)

Bases: TextGenerationModel

Generates completions using models deployed on inference servers like TRT or VLLM. This model assumes the default APIs are being used (e.g., no OpenAI-compatible API in VLLM)

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

parse_response(response, prompt_len)

Get the completion from a response, removing the prompt from it if required.

Parameters:

• response (Dict) – response from the inference server.

• prompt_len (int) – len of the prompt.

Returns:

completion.

Return type:

str

prepare_data(prompt, generation_config)

Prepares the json data of a request to the inference server.

Parameters:

• prompt (str) – a prompt.

• generation_config (Dict[str, Any]) – a generation config.

Returns:

json data for a request.

Return type:

Dict[str, Any]

class text_machina.src.models.openai.OpenAIModel(model_config)

Bases: TextGenerationModel

Generates completions using OpenAI models.

Requires the definition of the OPENAI_API_KEY=<key> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.vertex.VertexModel(model_config)

Bases: TextGenerationModel

Generates completions using VertexAI models. Requires the definition of the VERTEX_AI_CREDENTIALS_FILE=<path> environment variable.

generate_completion(prompt, generation_config)

Generates a completion for a prompt by decoding a model parameterized by generation_config. This method has to be overwritten to implement the completion code.

Parameters:

• prompts (str) – prompt to generate completions for.

• generation_config (Dict) – Dictionary containing the generation parameters.

Returns:

Generated completion or .types.GENERATION_ERROR if there was some error.

Return type:

str

class text_machina.src.models.types.CompletionType(value)

Bases: str, Enum

An enumeration.

CHAT: str = 'CHAT'

COMPLETION: str = 'COMPLETION'