Module `lore.types.base_types`

Expand source code

import datetime
import json
import logging
import re
import typing
from typing import Any, Dict, List, Optional, Type, Union

import datasets
from determined.common import api
from piccolo_api.crud.serializers import create_pydantic_model
from pydantic import BaseModel, Field

import lore.types.utils as utils
from lore import const
from lore.const import DEFAULT_SERVER_PORT
from lore.orm.lore_orm import tables as orm
from lore.types.enums import ControllerType, DatasetDataType, GPUType, ModelSource


class Workspace(BaseModel):
    id: int
    name: str
    experiment_project_id: int
    controller_project_id: int

    @staticmethod
    def _from_orm(workspace: orm.Workspace) -> "Workspace":
        return Workspace(
            id=workspace.id,
            name=workspace.name,
            experiment_project_id=workspace.experiment_project_id,
            controller_project_id=workspace.controller_project_id,
        )

    async def _to_orm(self) -> orm.Workspace:
        orm_workspace: orm.Workspace = None
        if self.id:
            orm_workspace = await orm.Workspace.objects().get(orm.Workspace.id == self.id)  # type: ignore[arg-type]
        if orm_workspace is None:
            orm_workspace = orm.Workspace(id=self.id)
        orm_workspace.name = self.name
        orm_workspace.experiment_project_id = self.experiment_project_id
        orm_workspace.controller_project_id = self.controller_project_id
        return orm_workspace


class CreateWorkspaceRequest(BaseModel):
    name: str


class User(BaseModel):
    username: str
    admin: bool
    displayName: Optional[str] = None
    remote: Optional[bool] = None


class LoreInfo(BaseModel):
    commit: str = Field()
    det_master: str = Field()


class Project(BaseModel):
    id: int = Field(default=0)
    name: str = Field(default="")
    workspace_id: int = Field(default=0)


class AsyncResponse(BaseModel):
    task_id: str = Field(default="")
    status: str = Field(default="")
    ready: bool = Field(default=False)
    successful: bool = Field(default=False)
    result: typing.Optional[Any] = Field(default="")
    error: typing.Optional[str] = Field(default="")


class LoreSession(BaseModel):
    workspace: Workspace = Field(...)  # "..." means required
    project: Project = Field(...)  # "..." means required
    lore_server_ip: str = "localhost"
    lore_server_prefix: str = "/lore"
    lore_server_port: int = DEFAULT_SERVER_PORT
    lore_server_root_path: str = ""
    lore_server_redis_port: int = 9013


class runtimeStreamError(BaseModel):
    message: str = Field(default="")


class UploadedFiles(BaseModel):
    file_paths: List[str] = Field(default=[])


_Dataset: Type[BaseModel] = create_pydantic_model(table=orm.Dataset, model_name="_Dataset")


class Dataset(_Dataset):  # type: ignore[valid-type, misc]
    registered_in_genai: bool = True
    id: Optional[int] = None


class ConstructDatasetFromHFRequest(BaseModel):
    dataset_name: str
    workspace_id: int
    dataset_config_name: Optional[str] = None
    token: Optional[str] = None
    task_type: Optional[str] = None
    data_type: Optional[List[DatasetDataType]] = None
    source: Optional[str] = None


class ConstructDatasetFromLocalRequest(BaseModel):
    dataset_name: str
    workspace_id: int
    # File containing all the arrow files. Needs to be zip
    arrow_file_path: Optional[str]
    data_type: Optional[List[DatasetDataType]] = None
    # train_files/validation_files/test_files are filepaths
    # (relative to uploaded_files_uuid shared_fs directory)
    # that indicate which files should be used to build train/validation/test split
    # in a dataset. File names can contain regexes.
    train_files: Optional[List[str]] = None
    validation_files: Optional[List[str]] = None
    test_files: Optional[List[str]] = None


class MergeAndResplitRequest(BaseModel):
    dataset: Optional[Dataset] = None
    dataset_id: Optional[int] = None
    train_ratio: Optional[float] = None
    test_ratio: Optional[float] = None
    validation_ratio: Optional[float] = None
    splits_to_resplit: List[str] = None
    workspace_id: int
    shuffle: bool = False
    seed: Optional[int] = None


class SampleDatasetRequest(BaseModel):
    dataset: Optional[Dataset] = None
    dataset_id: Optional[int] = None
    start_index: Optional[int] = None
    number_of_samples: Optional[int] = None
    ratio: Optional[float] = None
    splits: Optional[List[str]] = None
    seed: Optional[int] = None
    as_dataset: bool = False
    workspace_id: int


class StartEndTokens(BaseModel):
    start: str = ""
    end: str = ""


class FormattingTokens(BaseModel):
    input_role_tokens: Dict[str, StartEndTokens] = Field(
        default_factory=lambda: {
            "system": StartEndTokens(start="", end="\n"),
            "user": StartEndTokens(start="", end="\n"),
        }
    )
    output_role_tokens: Dict[str, StartEndTokens] = Field(
        default_factory=lambda: {
            "assistant": StartEndTokens(start="", end=""),
        }
    )
    input_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")
    output_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")

    def _process_section(
        self,
        template: Union[str, Dict[str, str]],
        role_tokens: Dict[str, StartEndTokens],
        role_ordering: List[str],
        sequence_tokens: StartEndTokens,
        append_special_tokens: bool = True,
    ) -> str:
        template_str = ""

        if isinstance(template, dict):
            formatted_parts = [
                (
                    f"{role_tokens[role].start}{template[role].strip()}{role_tokens[role].end}"
                    if append_special_tokens
                    else f"{template[role].strip()}"
                )
                for role in role_ordering
                if role in template
            ]
            template_str = "".join(formatted_parts)

        else:
            template_str = template.strip()

        if template_str:
            # Append start and end tokens only if template_str is not empty
            prefix = sequence_tokens.start if append_special_tokens else ""
            suffix = sequence_tokens.end if append_special_tokens else ""
            template_str = f"{prefix}{template_str}{suffix}"

        return template_str

    def _wrap(
        self,
        template: Union[str, Dict[str, str]],
        in_section: bool = True,
        out_section: bool = True,
        append_special_tokens: bool = True,
    ) -> str:
        if in_section:
            in_section_text = self._process_section(
                template,
                self.input_role_tokens,
                self._input_role_ordering(),
                self.input_sequence_tokens,
            )
        else:
            in_section_text = ""

        # Process output section only if PromptTemplate is defined as a dictionary.
        # If a prompt template is a str, then it is assumed to contain only user input.
        if out_section and isinstance(template, dict):
            out_section_text = self._process_section(
                template,
                self.output_role_tokens,
                self._output_role_ordering(),
                self.output_sequence_tokens,
                append_special_tokens=append_special_tokens,
            )
        else:
            out_section_text = ""

        prompt = f"{in_section_text}{out_section_text}"
        if len(out_section_text.strip()) == 0:
            # Append a role token to continue text generation if 'out_section_text' is empty.
            # Select the 'assistant' start token or default to the first token if 'assistant' is undefined.
            assistant_key = (
                "assistant"
                if "assistant" in self.output_role_tokens
                else next(iter(self.output_role_tokens))
            )
            prompt += self.output_role_tokens[assistant_key].start
        return prompt

    def _input_role_ordering(self) -> List[str]:
        ordering = ["system"] if "system" in self.input_role_tokens else []
        ordering += ["user"] if "user" in self.input_role_tokens else []
        ordering += [key for key, _ in self.input_role_tokens.items() if key not in ordering]
        return ordering

    def _output_role_ordering(self) -> List[str]:
        ordering = ["assistant"] if "assistant" in self.output_role_tokens else []
        return ordering + [key for key, _ in self.output_role_tokens.items() if key not in ordering]

    def __str__(self) -> str:
        template = {key: f"{{{key}}}" for key in self._input_role_ordering()}
        template.update({key: f"{{{key}}}" for key in self._output_role_ordering()})
        return self._wrap(template)


class DefaultFormatting(FormattingTokens):
    input_role_tokens: Dict[str, StartEndTokens] = {
        "system": StartEndTokens(start="", end=""),
        "user": StartEndTokens(start="", end=""),
    }
    output_role_tokens: Dict[str, StartEndTokens] = {"assistant": StartEndTokens(start="", end="")}
    input_sequence_tokens: StartEndTokens = StartEndTokens(start="", end=" ")
    output_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")

    def _wrap(
        self,
        template: Union[str, Dict[str, str]],
        in_section: bool = True,
        out_section: bool = True,
        append_special_tokens: bool = True,
    ) -> str:
        if isinstance(template, str):
            return template

        in_section_text = (
            "\n".join(
                [f"{text}" for key, text in template.items() if key not in self.output_role_tokens]
            )
            if in_section
            else ""
        )
        out_section_text = (
            "\n".join(
                [f"{text}" for key, text in template.items() if key in self.output_role_tokens]
            )
            if out_section
            else ""
        )

        return f"{in_section_text}{out_section_text}".strip()


# Model related types.


class PEFTConfig(BaseModel):
    # Add options from the `peft` library as desired.
    peft_type: Optional[str] = None
    """The variety of PEFT to use."""

    peft_args: Dict[str, Any] = Field(default_factory=lambda: {})
    """Arguments to pass to the PEFT constructor."""


class Model(BaseModel):
    genai_model_version: Optional[int] = None
    genai_model_name: Optional[str] = None
    genai_checkpoint_uuid: Optional[str] = None
    hf_model_name: Optional[str] = None
    hf_model_revision: Optional[str] = None
    hf_model_commit: Optional[str] = None
    model_architecture: str = ""
    problem_type: str = ""
    source: Optional[ModelSource] = None
    workspace_id: Optional[int] = None
    source_experiment_id: Optional[int] = None
    model_config: Optional[Dict[str, Any]] = None
    peft_config: Optional["PEFTConfig"] = None
    peft_checkpoint_uuid: Optional[str] = None
    formatting_tokens: FormattingTokens = None
    registered_in_genai: bool = False
    id: Optional[int] = None
    creation_time: Optional[datetime.datetime] = None
    modified_time: Optional[datetime.datetime] = None

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        self.registered_in_genai = self.id is not None

    def name(self) -> str:
        return self.genai_model_name if self.source == ModelSource.GENAI else self.hf_model_name

    @classmethod
    def from_hf(
        cls,
        hf_model_name: str,
        branch: Optional[str] = None,
        hf_model_commit: Optional[str] = None,
        **model_kwargs: Any,  # Passed to bt.Model constructor.
    ):
        if hf_model_commit is not None:
            if branch is not None:
                utils.check_commit_hash_branch_association(hf_model_name, branch, hf_model_commit)
            else:
                # Huggingface API does not provide an easy way to look up branch from commit_hash
                raise RuntimeError("If commit_hash is provided, branch must be provided as well")
        else:  # commit_hash is None
            if branch is None:
                branch = "main"
            hf_model_commit = utils.get_latest_commit_hash_from_branch(hf_model_name, branch)
            if hf_model_commit is None:
                raise RuntimeError(
                    f"Branch {branch} does not exist for model {hf_model_name}. Please provide a valid branch.",
                )

        # Rename fields
        # TODO: [MLG-1243]The revision field will be renamed branch later
        return cls(
            hf_model_name=hf_model_name,
            hf_model_commit=hf_model_commit,
            hf_model_revision=branch,
            source=ModelSource.HUGGING_FACE,
            **model_kwargs,
        )

    @classmethod
    def _from_orm(cls, orm_model: orm.Model) -> "Model":
        kwargs = orm_model.to_dict()
        if kwargs["model_config"]:
            kwargs["model_config"] = json.loads(kwargs["model_config"])
        if kwargs["formatting_tokens"]:
            kwargs["formatting_tokens"] = FormattingTokens(
                **json.loads(kwargs["formatting_tokens"])
            )
        return Model(**kwargs)

    def _to_orm(self) -> orm.Model:
        kwargs = self.dict()
        if kwargs["model_config"]:
            kwargs["model_config"] = json.dumps(kwargs["model_config"])
        if kwargs["formatting_tokens"]:
            kwargs["formatting_tokens"] = json.dumps(kwargs["formatting_tokens"])
        if kwargs["creation_time"] is None:
            kwargs["creation_time"] = datetime.datetime.now()
        if kwargs["modified_time"] is None:
            kwargs["modified_time"] = datetime.datetime.now()
        # Pop fields if needed
        if kwargs["id"] is None:
            kwargs.pop("id")

        # Fields we always pop (not orm fields)
        list(map(kwargs.pop, ["registered_in_genai", "peft_config", "peft_checkpoint_uuid"]))
        return orm.Model(**kwargs)

    def is_valid(self):
        # TODO: Add additional validation for allowed models and if HF models are available.
        if self.source == ModelSource.GENAI:
            if self.registered_in_genai:
                return self.genai_model_name is not None and self.genai_model_version is not None
            else:
                return self.genai_checkpoint_uuid is not None
        elif self.source == ModelSource.HUGGING_FACE:
            return self.hf_model_name is not None
        else:
            return False


class ModelLoadConfig(BaseModel):
    torch_dtype: Optional[str] = Field(
        default=None,
        description=(
            "Override the default `torch.dtype` and load the model under this dtype. "
            "If `auto` is passed, the dtype will be automatically derived from "
            "the model's weights. "
            "One of auto, bfloat16, float16, float32."
        ),
    )
    low_cpu_mem_usage: bool = Field(
        default=False,
        description=(
            "It is an option to create the model as an empty shell, then only materialize "
            "its parameters when the pretrained weights are loaded."
            "set True will benefit LLM loading time and RAM consumption."
        ),
    )
    trust_remote_code: bool = Field(
        default=False,
        description=("Set to true to allow execution of code from remote model repos."),
    )
    token: Optional[str] = Field(default=None, description=("Hugging Face authorization token."))

    vllm_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description=(
            "VLLM config to use for loading a model for inference. "
            "If not provided, the default model config will be used."
        ),
    )

    peft_config: Optional["PEFTConfig"] = Field(
        default=None,
        description=("PEFT config to use for loading a model for training."),
    )


class TrainingCheckpointUserData(BaseModel):
    """
    Custom data stored in checkpoint directories.
    """

    base_model_name: str
    """
    The Huggingface repo on which this model is based.  Should not be a local path.
    """

    base_model_revision: str = "main"
    """
    The specific model version (can be a branch name, tag name or commit id).
    """

    base_model_uuid: Optional[str] = None
    """
    The checkpoint UUID this model was based on, if any.  Should not be the UUID of the checkpoint
    itself, but instead of its parent model.
    """

    peft_config: Optional[PEFTConfig] = None
    """
    The PEFT configuration used to train this model, if any.
    """


_Experiment: Type[BaseModel] = create_pydantic_model(table=orm.Experiment, model_name="_Experiment")


class Experiment(_Experiment):  # type: ignore[valid-type, misc]
    registered_in_genai: bool = True
    state: str = ""
    datasets: List["Dataset"] = Field(default_factory=list)
    source_models: List["Model"] = Field(default_factory=list)

    @classmethod
    async def _aync_from_orm(cls, data: orm.Experiment) -> "Experiment":
        """TODO: Populate m2m relationships after API decision."""
        result = cls(**data.to_dict())
        return result

    def to_orm(self) -> orm.Experiment:
        kwargs = self.dict()
        kwargs.pop("registered_in_genai")
        kwargs.pop("state")
        kwargs.pop("datasets")
        kwargs.pop("source_models")
        return orm.Experiment(**kwargs)


class PromptTemplate(BaseModel):
    name: Optional[str] = None
    template: Union[str, Dict[str, str]] = ""
    workspace_id: Optional[int] = None
    creation_time: Optional[datetime.datetime] = None
    modified_time: Optional[datetime.datetime] = None
    start_prompt_template: bool = False
    registered_in_genai: bool = True
    id: Optional[int] = None
    system_generated: bool = False

    @classmethod
    def _from_orm(cls, orm_prompt_template: orm.PromptTemplate) -> "PromptTemplate":
        kwargs = orm_prompt_template.to_dict()
        if kwargs["template"]:
            kwargs["template"] = json.loads(kwargs["template"])
            if len(kwargs["template"]) == 1 and "<<<genai_template>>>" in kwargs["template"]:
                kwargs["template"] = kwargs["template"]["<<<genai_template>>>"]
        return PromptTemplate(**kwargs)

    @classmethod
    def _to_orm(cls, prompt_template: "PromptTemplate") -> orm.PromptTemplate:
        kwargs = dict(prompt_template)
        kwargs.pop("id")
        kwargs.pop("registered_in_genai")
        if isinstance(prompt_template.template, dict):
            kwargs["template"] = json.dumps(prompt_template.template)
        else:
            kwargs["template"] = json.dumps({"<<<genai_template>>>": prompt_template.template})
        return orm.PromptTemplate(**kwargs)

    def _format_dataset_column(
        self,
        dataset: Union[datasets.Dataset, datasets.DatasetDict],
        template_input: str,
        column_name: str,
        batch_size: int = 10,
    ):
        feature_substitutions = set(re.findall(r"\{\{([a-zA-Z0-9_\- ]+)\}\}", template_input))

        def apply_prompts(examples: Dict[str, List]):
            vals = []
            batch_length = min(batch_size, len(examples[list(examples.keys())[0]]))
            for example_idx in range(batch_length):
                args = {
                    feature_name: examples[feature_name][example_idx]
                    for feature_name in examples
                    if feature_name in feature_substitutions
                }
                formatted_template = str(template_input)
                for key, value in args.items():
                    formatted_template = formatted_template.replace("{{" + key + "}}", str(value))
                vals.append(formatted_template)
            examples[column_name] = vals
            return examples

        return dataset.map(apply_prompts, batched=True, batch_size=batch_size)

    def _format_dataset(
        self,
        dataset: Union[datasets.Dataset, datasets.DatasetDict],
        formatting: FormattingTokens,
        new_column_name: str = "genai_input_text",
        expected_output_column_name: Optional[str] = None,
        batch_size: int = 10,
    ):
        if expected_output_column_name:
            # If expected output_column_name is provided, put system and user content in
            # new_column_name, and assistant content in expected_output_column_name.
            template_input = formatting._wrap(
                template=self.template, in_section=True, out_section=False
            )
            dataset = self._format_dataset_column(
                dataset, template_input, new_column_name, batch_size
            )

            template_output = formatting._wrap(
                template=self.template,
                in_section=False,
                out_section=True,
                append_special_tokens=False,
            )
            dataset = self._format_dataset_column(
                dataset, template_output, expected_output_column_name, batch_size
            )
        else:
            # If expected output_column_name is not provided, put all content in
            # new_column_name
            template_input = formatting._wrap(
                template=self.template, in_section=True, out_section=True
            )
            dataset = self._format_dataset_column(
                dataset, template_input, new_column_name, batch_size
            )
        return dataset

    def format(
        self,
        input: Optional[Union[Dict, datasets.Dataset, datasets.DatasetDict]] = None,
        model: Optional[Model] = None,
        formatting_tokens: Optional[FormattingTokens] = None,
        **kwargs: Any,
    ) -> Union[str, datasets.Dataset, datasets.DatasetDict]:
        if formatting_tokens is not None and model is not None:
            raise ValueError("Exactly one of model or formatting can be provided.")
        elif formatting_tokens is None and model is None:
            logging.warning("Using default input formatting.")
            formatting = DefaultFormatting()  # type: FormattingTokens
        else:
            formatting = model.formatting_tokens if model else formatting_tokens

        template = formatting._wrap(self.template)
        if input is None:
            return template
        elif isinstance(input, datasets.Dataset) or isinstance(input, datasets.DatasetDict):
            return self._format_dataset(input, formatting, **kwargs)
        else:
            for key, value in input.items():
                template = template.replace("{{" + key + "}}", str(value))
            return template


class PlaygroundSnapshot(BaseModel):
    name: str
    id: Optional[int] = None
    project_id: int
    model_id: int
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    modified_time: Optional[datetime.datetime] = None
    web: Optional[Dict[str, Any]] = None

    @classmethod
    def _from_orm(cls, orm_playground_snapshot: orm.PlaygroundSnapshot) -> "PlaygroundSnapshot":
        kwargs = orm_playground_snapshot.to_dict()
        if kwargs["generation_config"]:
            kwargs["generation_config"] = json.loads(kwargs["generation_config"])
        if kwargs["prompt"]:
            kwargs["prompt"] = json.loads(kwargs["prompt"])
        if kwargs["dataset_samples"]:
            kwargs["dataset_samples"] = json.loads(kwargs["dataset_samples"])
        if kwargs["web"]:
            kwargs["web"] = json.loads(kwargs["web"])

        return PlaygroundSnapshot.parse_obj(kwargs)


class DataArguments(BaseModel):
    dataset_data: Optional[Dict[str, Any]] = None
    block_size: Optional[int] = None


class TrainingConfig(BaseModel):
    data_args: DataArguments = Field(default_factory=DataArguments)
    slots_per_trial: int = 1
    max_steps: int = -1
    num_train_epochs: Optional[float] = None
    save_total_limit: int = 1
    learning_rate: float = 5e-5
    per_device_train_batch_size: int = const.DEFAULT_PER_DEVICE_TRAIN_BATCH_SIZE
    per_device_eval_batch_size: int = const.DEFAULT_PER_DEVICE_EVAL_BATCH_SIZE
    do_train: bool = True
    do_eval: bool = True
    fp16: bool = False  # Deepspeed flag
    bf16: bool = False  # Deepspeed flag
    gradient_checkpointing: bool = const.DEFAULT_TRAIN_GRADIENT_CHECKPOINTING_FLAG
    deepspeed: bool = const.DEFAULT_TRAIN_DEEPSPEED_FLAG
    ds_cpu_offload: bool = const.DEFAULT_TRAIN_DS_CPU_OFF_LOAD
    resource_pool: Optional[str] = None
    seed: int = 1337
    output_dir: str = "/tmp/lore_output"
    eval_steps: Optional[int] = Field(default=None)
    logging_steps: Optional[int] = Field(default=None)
    save_steps: Optional[int] = Field(default=None)
    logging_strategy: str = Field(default="")
    evaluation_strategy: str = Field(default="")
    save_strategy: str = Field(default="")
    save_on_each_node: bool = True

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        if self.logging_strategy == "":
            self.logging_strategy = "steps"  # The HF default.

        if self.max_steps != -1 and self.num_train_epochs is not None:
            raise ValueError("Only one of `max_steps` or `num_train_epochs` can be specified.")

        if self.max_steps != -1:
            self.evaluation_strategy = (
                self.evaluation_strategy if self.evaluation_strategy != "" else "steps"
            )
            self.eval_steps = (
                self.eval_steps
                if self.eval_steps
                else max(self.max_steps // self.save_total_limit, 1)
            )

            self.save_strategy = self.save_strategy if self.save_strategy != "" else "steps"
            self.save_steps = (
                self.save_steps
                if self.save_steps
                else max(self.max_steps // self.save_total_limit, 1)
            )
            self.logging_steps = (
                self.logging_steps
                if self.logging_steps
                else max(min(500, self.max_steps // self.save_total_limit), 1)
            )
        else:
            self.num_train_epochs = self.num_train_epochs if self.num_train_epochs else 3.0

            self.evaluation_strategy = (
                self.evaluation_strategy if self.evaluation_strategy != "" else "epoch"
            )
            self.eval_steps = (
                self.eval_steps
                if self.eval_steps
                else max(self.num_train_epochs // self.save_total_limit, 1)
            )

            self.save_strategy = self.save_strategy if self.save_strategy != "" else "epoch"
            self.save_steps = (
                self.save_steps
                if self.save_steps
                else max(self.num_train_epochs // self.save_total_limit, 1)
            )

            self.logging_steps = (
                self.logging_steps if self.logging_steps else 500
            )  # The HF default.


class LoreEvalTaskConfig(BaseModel):
    dataset: Optional[Dataset] = None


class LoreMCEvalTaskConfig(LoreEvalTaskConfig):
    query_col: str = "query"
    choices_col: str = "choices"
    gold_col: str = "gold"


class EvaluationConfig(BaseModel):
    data_args: DataArguments = Field(default_factory=DataArguments)
    tasks: Union[str, List[Union[str, Dict[str, Any]]]] = Field(
        default_factory=lambda: ["arc_easy"]
    )
    num_fewshot: int = Field(default=0)
    use_accelerate: bool = Field(default=False)
    slots_per_trial: int = Field(default=1)
    max_concurrent_trials: int = Field(default=1)

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        if isinstance(self.tasks, str):
            self.tasks = [self.tasks]


class ExperimentSettings(BaseModel):
    model: Model
    project_id: int
    name: Optional[str] = None
    train_config: Optional[TrainingConfig] = None
    eval_config: Optional[EvaluationConfig] = None
    model_load_config: ModelLoadConfig = ModelLoadConfig()
    prompt_template: Optional[PromptTemplate] = None
    resource_pool: Optional[str] = None

    def dataset(self) -> Optional[Dataset]:
        dataset: Optional[Dataset] = None
        if self.train_config:
            if self.train_config.data_args:
                if self.train_config.data_args.dataset_data:
                    dataset = Dataset.parse_obj(self.train_config.data_args.dataset_data)
        elif self.eval_config:
            if self.eval_config.data_args:
                if self.eval_config.data_args.dataset_data:
                    dataset = Dataset.parse_obj(self.eval_config.data_args.dataset_data)
        return dataset


class GenerateChatResponseRequest(BaseModel):
    prompts: Union[str, List[str]]
    generation_config: Dict[str, Any] = Field(default_factory=dict)
    batch_size: Optional[int] = None
    client_key: str = ""


class GenerateOnDatasetRequest(BaseModel):
    dataset: Dataset
    prompt_template: PromptTemplate
    output_column_name: str = "genai_generated_text"
    split: Optional[str] = None
    start_index: Optional[int] = None
    number_of_samples: Optional[int] = None
    expected_output_column_name: Optional[str] = "genai_expected_text"
    batch_size: int = 10
    test_mode: Optional[bool] = False
    generation_config: Dict[str, Any] = Field(default_factory=dict)
    num_proc: Optional[int] = None
    client_key: str = ""


class CreatePlaygroundSnaphotRequest(BaseModel):
    name: str
    project_id: int
    model_id: int
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    output_datasets: Optional[List[Dataset]] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    web: Optional[Dict[str, Any]] = None


class UpdatePlaygroundSnaphotRequest(BaseModel):
    # Snapshot id will be from URI, other fields are all optional as user can choose to update
    # any number of them
    id: Optional[int] = None
    name: Optional[str] = None
    model_id: Optional[int] = None
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    output_datasets: Optional[List[Dataset]] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    web: Optional[Dict[str, Any]] = None


class ApplyPromptTemplateRequest(BaseModel):
    dataset: Dataset
    prompt_template: PromptTemplate
    workspace_id: int
    model: Optional[Model] = None
    new_column_name: str = "genai_input_text"
    expected_output_column_name: Optional[str] = None
    batch_size: int = 10


class ConcatenateDatasetsRequest(BaseModel):
    datasets: List[Dataset]
    name: str
    workspace_id: int
    system_generated: bool = False
    register_dataset: bool = False


class ComputeMetricsRequest(BaseModel):
    dataset: Dataset
    workspace_id: int
    ground_truth_column_name: str
    predictions_column_name: str
    metrics: List[str] = ["exact_match"]
    substr_match: bool = True
    split: Optional[str] = None
    strip: bool = False
    lower: bool = False


class RestartControllerRequest(BaseModel):
    controller_type: ControllerType
    workspace_id: int

    # Currently user's ip (can be located in experiment hyperparameters celery_queue).
    # If not specified, the requester's unique id (ip address) will be used.
    # TODO: Change to username later.
    unique_id: Optional[str] = None
    client_key: Optional[str] = None


class RestartControllerResponse(BaseModel):
    success: bool = None
    prev_experiment_id: Optional[int] = None
    experiment_id: Optional[int] = None


class RegisterHFModelRequest(BaseModel):
    hf_model_name: str
    branch: Optional[str] = None
    hf_model_commit: Optional[str] = None
    formatting_tokens: Optional[FormattingTokens] = None
    problem_type: str = ""
    model_config: Optional[Dict[str, Any]] = None


class UploadModelToHFRequest(BaseModel):
    model: Model
    hf_repo_owner: str
    hf_repo_name: str
    hf_token: str
    private: bool


class UploadModelToHFResponse(BaseModel):
    experiment_id: int


class DeleteDatasetResponse(BaseModel):
    task_type: str
    task_status: str
    dataset_id: int
    error_message: Optional[str]


class ResourcePoolMetadata(BaseModel):
    max_agents: int
    gpu_type: GPUType


class ResourcePool(BaseModel):
    name: str
    gpu_type: Optional[GPUType] = None
    slots_per_agent: int
    slots_available: int
    slots_used: int
    max_agents: int

    @classmethod
    def from_bindings(
        cls,
        bindings: api.bindings.v1ResourcePool,
        extra_metadata: Dict[str, ResourcePoolMetadata] = {},
    ) -> "ResourcePool":
        name = bindings.name
        if name in extra_metadata:
            gpu_type = extra_metadata[name].gpu_type
            max_agents = extra_metadata[name].max_agents
        else:
            gpu_type = GPUType.from_string(bindings.accelerator)
            max_agents = bindings.maxAgents
        slots_per_agent = bindings.slotsPerAgent
        slots_available = bindings.slotsAvailable
        slots_used = bindings.slotsUsed
        return ResourcePool(
            name=name,
            gpu_type=gpu_type,
            slots_per_agent=slots_per_agent,
            slots_available=slots_available,
            slots_used=slots_used,
            max_agents=max_agents,
        )


class VLLMInferenceConfig(BaseModel):
    slots_per_trial: int
    max_new_tokens: int
    batch_size: int
    swap_space: int
    torch_dtype: str


class InferenceConfigRequest(BaseModel):
    model_id: int
    # Max number of resource pool to return
    max_resource_pool: Optional[int] = None
    # Max number of config to return per resource pool
    max_config_per_resource_pool: Optional[int] = None


class ResourcePoolInferenceConfigs(BaseModel):
    resource_pool: ResourcePool
    vllm_configs: List[VLLMInferenceConfig]


class GetModelByNameRequest(BaseModel):
    workspace_id: int
    name: str
    is_base_model: bool
    genai_model_version: Optional[int]


class TrainingConfigRequest(BaseModel):
    model_id: Optional[int]
    model_architecture: Optional[str]
    # Max number of resource pool to return
    max_resource_pool: Optional[int] = None
    # Max number of config to return per resource pool
    max_config_per_resource_pool: Optional[int] = None


class TrainingConfigRecommendation(BaseModel):
    slots_per_trial: int
    context_window: int
    batch_size: int
    deepspeed: bool
    gradient_checkpointing: bool
    torch_dtype: str


class ResourcePoolTrainingConfigs(BaseModel):
    resource_pool: ResourcePool
    training_configs: List[TrainingConfigRecommendation]

Classes

class ApplyPromptTemplateRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ApplyPromptTemplateRequest(BaseModel):
    dataset: Dataset
    prompt_template: PromptTemplate
    workspace_id: int
    model: Optional[Model] = None
    new_column_name: str = "genai_input_text"
    expected_output_column_name: Optional[str] = None
    batch_size: int = 10

Ancestors

pydantic.main.BaseModel

Class variables

var batch_size : int
var dataset : Dataset
var expected_output_column_name : Optional[str]
var model : Optional[Model]
var model_computed_fields
var model_config
var model_fields
var new_column_name : str
var prompt_template : PromptTemplate
var workspace_id : int

class AsyncResponse (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class AsyncResponse(BaseModel):
    task_id: str = Field(default="")
    status: str = Field(default="")
    ready: bool = Field(default=False)
    successful: bool = Field(default=False)
    result: typing.Optional[Any] = Field(default="")
    error: typing.Optional[str] = Field(default="")

Ancestors

pydantic.main.BaseModel

Class variables

var error : Optional[str]
var model_computed_fields
var model_config
var model_fields
var ready : bool
var result : Optional[Any]
var status : str
var successful : bool
var task_id : str

class ComputeMetricsRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ComputeMetricsRequest(BaseModel):
    dataset: Dataset
    workspace_id: int
    ground_truth_column_name: str
    predictions_column_name: str
    metrics: List[str] = ["exact_match"]
    substr_match: bool = True
    split: Optional[str] = None
    strip: bool = False
    lower: bool = False

Ancestors

pydantic.main.BaseModel

Class variables

var dataset : Dataset
var ground_truth_column_name : str
var lower : bool
var metrics : List[str]
var model_computed_fields
var model_config
var model_fields
var predictions_column_name : str
var split : Optional[str]
var strip : bool
var substr_match : bool
var workspace_id : int

class ConcatenateDatasetsRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ConcatenateDatasetsRequest(BaseModel):
    datasets: List[Dataset]
    name: str
    workspace_id: int
    system_generated: bool = False
    register_dataset: bool = False

Ancestors

pydantic.main.BaseModel

Class variables

var datasets : List[Dataset]
var model_computed_fields
var model_config
var model_fields
var name : str
var register_dataset : bool
var system_generated : bool
var workspace_id : int

class ConstructDatasetFromHFRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ConstructDatasetFromHFRequest(BaseModel):
    dataset_name: str
    workspace_id: int
    dataset_config_name: Optional[str] = None
    token: Optional[str] = None
    task_type: Optional[str] = None
    data_type: Optional[List[DatasetDataType]] = None
    source: Optional[str] = None

Ancestors

pydantic.main.BaseModel

Class variables

var data_type : Optional[List[DatasetDataType]]
var dataset_config_name : Optional[str]
var dataset_name : str
var model_computed_fields
var model_config
var model_fields
var source : Optional[str]
var task_type : Optional[str]
var token : Optional[str]
var workspace_id : int

class ConstructDatasetFromLocalRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ConstructDatasetFromLocalRequest(BaseModel):
    dataset_name: str
    workspace_id: int
    # File containing all the arrow files. Needs to be zip
    arrow_file_path: Optional[str]
    data_type: Optional[List[DatasetDataType]] = None
    # train_files/validation_files/test_files are filepaths
    # (relative to uploaded_files_uuid shared_fs directory)
    # that indicate which files should be used to build train/validation/test split
    # in a dataset. File names can contain regexes.
    train_files: Optional[List[str]] = None
    validation_files: Optional[List[str]] = None
    test_files: Optional[List[str]] = None

Ancestors

pydantic.main.BaseModel

Class variables

var arrow_file_path : Optional[str]
var data_type : Optional[List[DatasetDataType]]
var dataset_name : str
var model_computed_fields
var model_config
var model_fields
var test_files : Optional[List[str]]
var train_files : Optional[List[str]]
var validation_files : Optional[List[str]]
var workspace_id : int

class CreatePlaygroundSnaphotRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class CreatePlaygroundSnaphotRequest(BaseModel):
    name: str
    project_id: int
    model_id: int
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    output_datasets: Optional[List[Dataset]] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    web: Optional[Dict[str, Any]] = None

Ancestors

pydantic.main.BaseModel

Class variables

var dataset_samples : Optional[Dict[str, Any]]
var generation_config : Optional[Dict[str, Any]]
var input_dataset_id : Optional[int]
var input_text : Optional[str]
var model_computed_fields
var model_config
var model_fields
var model_id : int
var name : str
var output_dataset_id : Optional[int]
var output_datasets : Optional[List[Dataset]]
var output_text : Optional[str]
var project_id : int
var prompt : Optional[Dict[str, Any]]
var starting_prompt_template_id : Optional[int]
var web : Optional[Dict[str, Any]]

class CreateWorkspaceRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class CreateWorkspaceRequest(BaseModel):
    name: str

Ancestors

pydantic.main.BaseModel

Class variables

var model_computed_fields
var model_config
var model_fields
var name : str

class DataArguments (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class DataArguments(BaseModel):
    dataset_data: Optional[Dict[str, Any]] = None
    block_size: Optional[int] = None

Ancestors

pydantic.main.BaseModel

Class variables

var block_size : Optional[int]
var dataset_data : Optional[Dict[str, Any]]
var model_computed_fields
var model_config
var model_fields

class Dataset (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Dataset(_Dataset):  # type: ignore[valid-type, misc]
    registered_in_genai: bool = True
    id: Optional[int] = None

Ancestors

piccolo.utils.pydantic._Dataset
pydantic.main.BaseModel

Class variables

var id : Optional[int]
var model_computed_fields
var model_config
var model_fields
var registered_in_genai : bool

class DefaultFormatting (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class DefaultFormatting(FormattingTokens):
    input_role_tokens: Dict[str, StartEndTokens] = {
        "system": StartEndTokens(start="", end=""),
        "user": StartEndTokens(start="", end=""),
    }
    output_role_tokens: Dict[str, StartEndTokens] = {"assistant": StartEndTokens(start="", end="")}
    input_sequence_tokens: StartEndTokens = StartEndTokens(start="", end=" ")
    output_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")

    def _wrap(
        self,
        template: Union[str, Dict[str, str]],
        in_section: bool = True,
        out_section: bool = True,
        append_special_tokens: bool = True,
    ) -> str:
        if isinstance(template, str):
            return template

        in_section_text = (
            "\n".join(
                [f"{text}" for key, text in template.items() if key not in self.output_role_tokens]
            )
            if in_section
            else ""
        )
        out_section_text = (
            "\n".join(
                [f"{text}" for key, text in template.items() if key in self.output_role_tokens]
            )
            if out_section
            else ""
        )

        return f"{in_section_text}{out_section_text}".strip()

Ancestors

FormattingTokens
pydantic.main.BaseModel

Class variables

var input_role_tokens : Dict[str, StartEndTokens]
var input_sequence_tokens : StartEndTokens
var model_computed_fields
var model_config
var model_fields
var output_role_tokens : Dict[str, StartEndTokens]
var output_sequence_tokens : StartEndTokens

class DeleteDatasetResponse (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class DeleteDatasetResponse(BaseModel):
    task_type: str
    task_status: str
    dataset_id: int
    error_message: Optional[str]

Ancestors

pydantic.main.BaseModel

Class variables

var dataset_id : int
var error_message : Optional[str]
var model_computed_fields
var model_config
var model_fields
var task_status : str
var task_type : str

class EvaluationConfig (*args, **data)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class EvaluationConfig(BaseModel):
    data_args: DataArguments = Field(default_factory=DataArguments)
    tasks: Union[str, List[Union[str, Dict[str, Any]]]] = Field(
        default_factory=lambda: ["arc_easy"]
    )
    num_fewshot: int = Field(default=0)
    use_accelerate: bool = Field(default=False)
    slots_per_trial: int = Field(default=1)
    max_concurrent_trials: int = Field(default=1)

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        if isinstance(self.tasks, str):
            self.tasks = [self.tasks]

Ancestors

pydantic.main.BaseModel

Class variables

var data_args : DataArguments
var max_concurrent_trials : int
var model_computed_fields
var model_config
var model_fields
var num_fewshot : int
var slots_per_trial : int
var tasks : Union[str, List[Union[str, Dict[str, Any]]]]
var use_accelerate : bool

class Experiment (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Experiment(_Experiment):  # type: ignore[valid-type, misc]
    registered_in_genai: bool = True
    state: str = ""
    datasets: List["Dataset"] = Field(default_factory=list)
    source_models: List["Model"] = Field(default_factory=list)

    @classmethod
    async def _aync_from_orm(cls, data: orm.Experiment) -> "Experiment":
        """TODO: Populate m2m relationships after API decision."""
        result = cls(**data.to_dict())
        return result

    def to_orm(self) -> orm.Experiment:
        kwargs = self.dict()
        kwargs.pop("registered_in_genai")
        kwargs.pop("state")
        kwargs.pop("datasets")
        kwargs.pop("source_models")
        return orm.Experiment(**kwargs)

Ancestors

piccolo.utils.pydantic._Experiment
pydantic.main.BaseModel

Class variables

var datasets : List[Dataset]
var model_computed_fields
var model_config
var model_fields
var registered_in_genai : bool
var source_models : List[Model]
var state : str

Methods

def to_orm(self) ‑> Experiment

Expand source code

def to_orm(self) -> orm.Experiment:
    kwargs = self.dict()
    kwargs.pop("registered_in_genai")
    kwargs.pop("state")
    kwargs.pop("datasets")
    kwargs.pop("source_models")
    return orm.Experiment(**kwargs)

class ExperimentSettings (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ExperimentSettings(BaseModel):
    model: Model
    project_id: int
    name: Optional[str] = None
    train_config: Optional[TrainingConfig] = None
    eval_config: Optional[EvaluationConfig] = None
    model_load_config: ModelLoadConfig = ModelLoadConfig()
    prompt_template: Optional[PromptTemplate] = None
    resource_pool: Optional[str] = None

    def dataset(self) -> Optional[Dataset]:
        dataset: Optional[Dataset] = None
        if self.train_config:
            if self.train_config.data_args:
                if self.train_config.data_args.dataset_data:
                    dataset = Dataset.parse_obj(self.train_config.data_args.dataset_data)
        elif self.eval_config:
            if self.eval_config.data_args:
                if self.eval_config.data_args.dataset_data:
                    dataset = Dataset.parse_obj(self.eval_config.data_args.dataset_data)
        return dataset

Ancestors

pydantic.main.BaseModel

Class variables

var eval_config : Optional[EvaluationConfig]
var model : Model
var model_computed_fields
var model_config
var model_fields
var model_load_config : ModelLoadConfig
var name : Optional[str]
var project_id : int
var prompt_template : Optional[PromptTemplate]
var resource_pool : Optional[str]
var train_config : Optional[TrainingConfig]

Methods

def dataset(self) ‑> Optional[Dataset]

Expand source code

def dataset(self) -> Optional[Dataset]:
    dataset: Optional[Dataset] = None
    if self.train_config:
        if self.train_config.data_args:
            if self.train_config.data_args.dataset_data:
                dataset = Dataset.parse_obj(self.train_config.data_args.dataset_data)
    elif self.eval_config:
        if self.eval_config.data_args:
            if self.eval_config.data_args.dataset_data:
                dataset = Dataset.parse_obj(self.eval_config.data_args.dataset_data)
    return dataset

class FormattingTokens (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class FormattingTokens(BaseModel):
    input_role_tokens: Dict[str, StartEndTokens] = Field(
        default_factory=lambda: {
            "system": StartEndTokens(start="", end="\n"),
            "user": StartEndTokens(start="", end="\n"),
        }
    )
    output_role_tokens: Dict[str, StartEndTokens] = Field(
        default_factory=lambda: {
            "assistant": StartEndTokens(start="", end=""),
        }
    )
    input_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")
    output_sequence_tokens: StartEndTokens = StartEndTokens(start="", end="")

    def _process_section(
        self,
        template: Union[str, Dict[str, str]],
        role_tokens: Dict[str, StartEndTokens],
        role_ordering: List[str],
        sequence_tokens: StartEndTokens,
        append_special_tokens: bool = True,
    ) -> str:
        template_str = ""

        if isinstance(template, dict):
            formatted_parts = [
                (
                    f"{role_tokens[role].start}{template[role].strip()}{role_tokens[role].end}"
                    if append_special_tokens
                    else f"{template[role].strip()}"
                )
                for role in role_ordering
                if role in template
            ]
            template_str = "".join(formatted_parts)

        else:
            template_str = template.strip()

        if template_str:
            # Append start and end tokens only if template_str is not empty
            prefix = sequence_tokens.start if append_special_tokens else ""
            suffix = sequence_tokens.end if append_special_tokens else ""
            template_str = f"{prefix}{template_str}{suffix}"

        return template_str

    def _wrap(
        self,
        template: Union[str, Dict[str, str]],
        in_section: bool = True,
        out_section: bool = True,
        append_special_tokens: bool = True,
    ) -> str:
        if in_section:
            in_section_text = self._process_section(
                template,
                self.input_role_tokens,
                self._input_role_ordering(),
                self.input_sequence_tokens,
            )
        else:
            in_section_text = ""

        # Process output section only if PromptTemplate is defined as a dictionary.
        # If a prompt template is a str, then it is assumed to contain only user input.
        if out_section and isinstance(template, dict):
            out_section_text = self._process_section(
                template,
                self.output_role_tokens,
                self._output_role_ordering(),
                self.output_sequence_tokens,
                append_special_tokens=append_special_tokens,
            )
        else:
            out_section_text = ""

        prompt = f"{in_section_text}{out_section_text}"
        if len(out_section_text.strip()) == 0:
            # Append a role token to continue text generation if 'out_section_text' is empty.
            # Select the 'assistant' start token or default to the first token if 'assistant' is undefined.
            assistant_key = (
                "assistant"
                if "assistant" in self.output_role_tokens
                else next(iter(self.output_role_tokens))
            )
            prompt += self.output_role_tokens[assistant_key].start
        return prompt

    def _input_role_ordering(self) -> List[str]:
        ordering = ["system"] if "system" in self.input_role_tokens else []
        ordering += ["user"] if "user" in self.input_role_tokens else []
        ordering += [key for key, _ in self.input_role_tokens.items() if key not in ordering]
        return ordering

    def _output_role_ordering(self) -> List[str]:
        ordering = ["assistant"] if "assistant" in self.output_role_tokens else []
        return ordering + [key for key, _ in self.output_role_tokens.items() if key not in ordering]

    def __str__(self) -> str:
        template = {key: f"{{{key}}}" for key in self._input_role_ordering()}
        template.update({key: f"{{{key}}}" for key in self._output_role_ordering()})
        return self._wrap(template)

Ancestors

pydantic.main.BaseModel

Subclasses

DefaultFormatting

Class variables

var input_role_tokens : Dict[str, StartEndTokens]
var input_sequence_tokens : StartEndTokens
var model_computed_fields
var model_config
var model_fields
var output_role_tokens : Dict[str, StartEndTokens]
var output_sequence_tokens : StartEndTokens

class GenerateChatResponseRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class GenerateChatResponseRequest(BaseModel):
    prompts: Union[str, List[str]]
    generation_config: Dict[str, Any] = Field(default_factory=dict)
    batch_size: Optional[int] = None
    client_key: str = ""

Ancestors

pydantic.main.BaseModel

Class variables

var batch_size : Optional[int]
var client_key : str
var generation_config : Dict[str, Any]
var model_computed_fields
var model_config
var model_fields
var prompts : Union[str, List[str]]

class GenerateOnDatasetRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class GenerateOnDatasetRequest(BaseModel):
    dataset: Dataset
    prompt_template: PromptTemplate
    output_column_name: str = "genai_generated_text"
    split: Optional[str] = None
    start_index: Optional[int] = None
    number_of_samples: Optional[int] = None
    expected_output_column_name: Optional[str] = "genai_expected_text"
    batch_size: int = 10
    test_mode: Optional[bool] = False
    generation_config: Dict[str, Any] = Field(default_factory=dict)
    num_proc: Optional[int] = None
    client_key: str = ""

Ancestors

pydantic.main.BaseModel

Class variables

var batch_size : int
var client_key : str
var dataset : Dataset
var expected_output_column_name : Optional[str]
var generation_config : Dict[str, Any]
var model_computed_fields
var model_config
var model_fields
var num_proc : Optional[int]
var number_of_samples : Optional[int]
var output_column_name : str
var prompt_template : PromptTemplate
var split : Optional[str]
var start_index : Optional[int]
var test_mode : Optional[bool]

class GetModelByNameRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class GetModelByNameRequest(BaseModel):
    workspace_id: int
    name: str
    is_base_model: bool
    genai_model_version: Optional[int]

Ancestors

pydantic.main.BaseModel

Class variables

var genai_model_version : Optional[int]
var is_base_model : bool
var model_computed_fields
var model_config
var model_fields
var name : str
var workspace_id : int

class InferenceConfigRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class InferenceConfigRequest(BaseModel):
    model_id: int
    # Max number of resource pool to return
    max_resource_pool: Optional[int] = None
    # Max number of config to return per resource pool
    max_config_per_resource_pool: Optional[int] = None

Ancestors

pydantic.main.BaseModel

Class variables

var max_config_per_resource_pool : Optional[int]
var max_resource_pool : Optional[int]
var model_computed_fields
var model_config
var model_fields
var model_id : int

class LoreEvalTaskConfig (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class LoreEvalTaskConfig(BaseModel):
    dataset: Optional[Dataset] = None

Ancestors

pydantic.main.BaseModel

Subclasses

LoreMCEvalTaskConfig

Class variables

var dataset : Optional[Dataset]
var model_computed_fields
var model_config
var model_fields

class LoreInfo (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class LoreInfo(BaseModel):
    commit: str = Field()
    det_master: str = Field()

Ancestors

pydantic.main.BaseModel

Class variables

var commit : str
var det_master : str
var model_computed_fields
var model_config
var model_fields

class LoreMCEvalTaskConfig (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class LoreMCEvalTaskConfig(LoreEvalTaskConfig):
    query_col: str = "query"
    choices_col: str = "choices"
    gold_col: str = "gold"

Ancestors

LoreEvalTaskConfig
pydantic.main.BaseModel

Class variables

var choices_col : str
var gold_col : str
var model_computed_fields
var model_config
var model_fields
var query_col : str

class LoreSession (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class LoreSession(BaseModel):
    workspace: Workspace = Field(...)  # "..." means required
    project: Project = Field(...)  # "..." means required
    lore_server_ip: str = "localhost"
    lore_server_prefix: str = "/lore"
    lore_server_port: int = DEFAULT_SERVER_PORT
    lore_server_root_path: str = ""
    lore_server_redis_port: int = 9013

Ancestors

pydantic.main.BaseModel

Class variables

var lore_server_ip : str
var lore_server_port : int
var lore_server_prefix : str
var lore_server_redis_port : int
var lore_server_root_path : str
var model_computed_fields
var model_config
var model_fields
var project : Project
var workspace : Workspace

class MergeAndResplitRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class MergeAndResplitRequest(BaseModel):
    dataset: Optional[Dataset] = None
    dataset_id: Optional[int] = None
    train_ratio: Optional[float] = None
    test_ratio: Optional[float] = None
    validation_ratio: Optional[float] = None
    splits_to_resplit: List[str] = None
    workspace_id: int
    shuffle: bool = False
    seed: Optional[int] = None

Ancestors

pydantic.main.BaseModel

Class variables

var dataset : Optional[Dataset]
var dataset_id : Optional[int]
var model_computed_fields
var model_config
var model_fields
var seed : Optional[int]
var shuffle : bool
var splits_to_resplit : List[str]
var test_ratio : Optional[float]
var train_ratio : Optional[float]
var validation_ratio : Optional[float]
var workspace_id : int

class Model (*args, **data)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Model(BaseModel):
    genai_model_version: Optional[int] = None
    genai_model_name: Optional[str] = None
    genai_checkpoint_uuid: Optional[str] = None
    hf_model_name: Optional[str] = None
    hf_model_revision: Optional[str] = None
    hf_model_commit: Optional[str] = None
    model_architecture: str = ""
    problem_type: str = ""
    source: Optional[ModelSource] = None
    workspace_id: Optional[int] = None
    source_experiment_id: Optional[int] = None
    model_config: Optional[Dict[str, Any]] = None
    peft_config: Optional["PEFTConfig"] = None
    peft_checkpoint_uuid: Optional[str] = None
    formatting_tokens: FormattingTokens = None
    registered_in_genai: bool = False
    id: Optional[int] = None
    creation_time: Optional[datetime.datetime] = None
    modified_time: Optional[datetime.datetime] = None

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        self.registered_in_genai = self.id is not None

    def name(self) -> str:
        return self.genai_model_name if self.source == ModelSource.GENAI else self.hf_model_name

    @classmethod
    def from_hf(
        cls,
        hf_model_name: str,
        branch: Optional[str] = None,
        hf_model_commit: Optional[str] = None,
        **model_kwargs: Any,  # Passed to bt.Model constructor.
    ):
        if hf_model_commit is not None:
            if branch is not None:
                utils.check_commit_hash_branch_association(hf_model_name, branch, hf_model_commit)
            else:
                # Huggingface API does not provide an easy way to look up branch from commit_hash
                raise RuntimeError("If commit_hash is provided, branch must be provided as well")
        else:  # commit_hash is None
            if branch is None:
                branch = "main"
            hf_model_commit = utils.get_latest_commit_hash_from_branch(hf_model_name, branch)
            if hf_model_commit is None:
                raise RuntimeError(
                    f"Branch {branch} does not exist for model {hf_model_name}. Please provide a valid branch.",
                )

        # Rename fields
        # TODO: [MLG-1243]The revision field will be renamed branch later
        return cls(
            hf_model_name=hf_model_name,
            hf_model_commit=hf_model_commit,
            hf_model_revision=branch,
            source=ModelSource.HUGGING_FACE,
            **model_kwargs,
        )

    @classmethod
    def _from_orm(cls, orm_model: orm.Model) -> "Model":
        kwargs = orm_model.to_dict()
        if kwargs["model_config"]:
            kwargs["model_config"] = json.loads(kwargs["model_config"])
        if kwargs["formatting_tokens"]:
            kwargs["formatting_tokens"] = FormattingTokens(
                **json.loads(kwargs["formatting_tokens"])
            )
        return Model(**kwargs)

    def _to_orm(self) -> orm.Model:
        kwargs = self.dict()
        if kwargs["model_config"]:
            kwargs["model_config"] = json.dumps(kwargs["model_config"])
        if kwargs["formatting_tokens"]:
            kwargs["formatting_tokens"] = json.dumps(kwargs["formatting_tokens"])
        if kwargs["creation_time"] is None:
            kwargs["creation_time"] = datetime.datetime.now()
        if kwargs["modified_time"] is None:
            kwargs["modified_time"] = datetime.datetime.now()
        # Pop fields if needed
        if kwargs["id"] is None:
            kwargs.pop("id")

        # Fields we always pop (not orm fields)
        list(map(kwargs.pop, ["registered_in_genai", "peft_config", "peft_checkpoint_uuid"]))
        return orm.Model(**kwargs)

    def is_valid(self):
        # TODO: Add additional validation for allowed models and if HF models are available.
        if self.source == ModelSource.GENAI:
            if self.registered_in_genai:
                return self.genai_model_name is not None and self.genai_model_version is not None
            else:
                return self.genai_checkpoint_uuid is not None
        elif self.source == ModelSource.HUGGING_FACE:
            return self.hf_model_name is not None
        else:
            return False

Ancestors

pydantic.main.BaseModel

Class variables

var creation_time : Optional[datetime.datetime]
var formatting_tokens : FormattingTokens
var genai_checkpoint_uuid : Optional[str]
var genai_model_name : Optional[str]
var genai_model_version : Optional[int]
var hf_model_commit : Optional[str]
var hf_model_name : Optional[str]
var hf_model_revision : Optional[str]
var id : Optional[int]
var model_architecture : str
var model_computed_fields
var model_config : Optional[Dict[str, Any]]
var model_fields
var modified_time : Optional[datetime.datetime]
var peft_checkpoint_uuid : Optional[str]
var peft_config : Optional[PEFTConfig]
var problem_type : str
var registered_in_genai : bool
var source : Optional[ModelSource]
var source_experiment_id : Optional[int]
var workspace_id : Optional[int]

Static methods

def from_hf(hf_model_name: str, branch: Optional[str] = None, hf_model_commit: Optional[str] = None, **model_kwargs: Any)

Expand source code

@classmethod
def from_hf(
    cls,
    hf_model_name: str,
    branch: Optional[str] = None,
    hf_model_commit: Optional[str] = None,
    **model_kwargs: Any,  # Passed to bt.Model constructor.
):
    if hf_model_commit is not None:
        if branch is not None:
            utils.check_commit_hash_branch_association(hf_model_name, branch, hf_model_commit)
        else:
            # Huggingface API does not provide an easy way to look up branch from commit_hash
            raise RuntimeError("If commit_hash is provided, branch must be provided as well")
    else:  # commit_hash is None
        if branch is None:
            branch = "main"
        hf_model_commit = utils.get_latest_commit_hash_from_branch(hf_model_name, branch)
        if hf_model_commit is None:
            raise RuntimeError(
                f"Branch {branch} does not exist for model {hf_model_name}. Please provide a valid branch.",
            )

    # Rename fields
    # TODO: [MLG-1243]The revision field will be renamed branch later
    return cls(
        hf_model_name=hf_model_name,
        hf_model_commit=hf_model_commit,
        hf_model_revision=branch,
        source=ModelSource.HUGGING_FACE,
        **model_kwargs,
    )

Methods

def is_valid(self)

Expand source code

def is_valid(self):
    # TODO: Add additional validation for allowed models and if HF models are available.
    if self.source == ModelSource.GENAI:
        if self.registered_in_genai:
            return self.genai_model_name is not None and self.genai_model_version is not None
        else:
            return self.genai_checkpoint_uuid is not None
    elif self.source == ModelSource.HUGGING_FACE:
        return self.hf_model_name is not None
    else:
        return False

def name(self) ‑> str

Expand source code

def name(self) -> str:
    return self.genai_model_name if self.source == ModelSource.GENAI else self.hf_model_name

class ModelLoadConfig (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ModelLoadConfig(BaseModel):
    torch_dtype: Optional[str] = Field(
        default=None,
        description=(
            "Override the default `torch.dtype` and load the model under this dtype. "
            "If `auto` is passed, the dtype will be automatically derived from "
            "the model's weights. "
            "One of auto, bfloat16, float16, float32."
        ),
    )
    low_cpu_mem_usage: bool = Field(
        default=False,
        description=(
            "It is an option to create the model as an empty shell, then only materialize "
            "its parameters when the pretrained weights are loaded."
            "set True will benefit LLM loading time and RAM consumption."
        ),
    )
    trust_remote_code: bool = Field(
        default=False,
        description=("Set to true to allow execution of code from remote model repos."),
    )
    token: Optional[str] = Field(default=None, description=("Hugging Face authorization token."))

    vllm_config: Optional[Dict[str, Any]] = Field(
        default=None,
        description=(
            "VLLM config to use for loading a model for inference. "
            "If not provided, the default model config will be used."
        ),
    )

    peft_config: Optional["PEFTConfig"] = Field(
        default=None,
        description=("PEFT config to use for loading a model for training."),
    )

Ancestors

pydantic.main.BaseModel

Class variables

var low_cpu_mem_usage : bool
var model_computed_fields
var model_config
var model_fields
var peft_config : Optional[PEFTConfig]
var token : Optional[str]
var torch_dtype : Optional[str]
var trust_remote_code : bool
var vllm_config : Optional[Dict[str, Any]]

class PEFTConfig (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class PEFTConfig(BaseModel):
    # Add options from the `peft` library as desired.
    peft_type: Optional[str] = None
    """The variety of PEFT to use."""

    peft_args: Dict[str, Any] = Field(default_factory=lambda: {})
    """Arguments to pass to the PEFT constructor."""

Ancestors

pydantic.main.BaseModel

Class variables

var model_computed_fields
var model_config
var model_fields
var peft_args : Dict[str, Any]: Arguments to pass to the PEFT constructor.
var peft_type : Optional[str]: The variety of PEFT to use.

class PlaygroundSnapshot (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class PlaygroundSnapshot(BaseModel):
    name: str
    id: Optional[int] = None
    project_id: int
    model_id: int
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    modified_time: Optional[datetime.datetime] = None
    web: Optional[Dict[str, Any]] = None

    @classmethod
    def _from_orm(cls, orm_playground_snapshot: orm.PlaygroundSnapshot) -> "PlaygroundSnapshot":
        kwargs = orm_playground_snapshot.to_dict()
        if kwargs["generation_config"]:
            kwargs["generation_config"] = json.loads(kwargs["generation_config"])
        if kwargs["prompt"]:
            kwargs["prompt"] = json.loads(kwargs["prompt"])
        if kwargs["dataset_samples"]:
            kwargs["dataset_samples"] = json.loads(kwargs["dataset_samples"])
        if kwargs["web"]:
            kwargs["web"] = json.loads(kwargs["web"])

        return PlaygroundSnapshot.parse_obj(kwargs)

Ancestors

pydantic.main.BaseModel

Class variables

var dataset_samples : Optional[Dict[str, Any]]
var generation_config : Optional[Dict[str, Any]]
var id : Optional[int]
var input_dataset_id : Optional[int]
var input_text : Optional[str]
var model_computed_fields
var model_config
var model_fields
var model_id : int
var modified_time : Optional[datetime.datetime]
var name : str
var output_dataset_id : Optional[int]
var output_text : Optional[str]
var project_id : int
var prompt : Optional[Dict[str, Any]]
var starting_prompt_template_id : Optional[int]
var web : Optional[Dict[str, Any]]

class Project (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Project(BaseModel):
    id: int = Field(default=0)
    name: str = Field(default="")
    workspace_id: int = Field(default=0)

Ancestors

pydantic.main.BaseModel

Class variables

var id : int
var model_computed_fields
var model_config
var model_fields
var name : str
var workspace_id : int

class PromptTemplate (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class PromptTemplate(BaseModel):
    name: Optional[str] = None
    template: Union[str, Dict[str, str]] = ""
    workspace_id: Optional[int] = None
    creation_time: Optional[datetime.datetime] = None
    modified_time: Optional[datetime.datetime] = None
    start_prompt_template: bool = False
    registered_in_genai: bool = True
    id: Optional[int] = None
    system_generated: bool = False

    @classmethod
    def _from_orm(cls, orm_prompt_template: orm.PromptTemplate) -> "PromptTemplate":
        kwargs = orm_prompt_template.to_dict()
        if kwargs["template"]:
            kwargs["template"] = json.loads(kwargs["template"])
            if len(kwargs["template"]) == 1 and "<<<genai_template>>>" in kwargs["template"]:
                kwargs["template"] = kwargs["template"]["<<<genai_template>>>"]
        return PromptTemplate(**kwargs)

    @classmethod
    def _to_orm(cls, prompt_template: "PromptTemplate") -> orm.PromptTemplate:
        kwargs = dict(prompt_template)
        kwargs.pop("id")
        kwargs.pop("registered_in_genai")
        if isinstance(prompt_template.template, dict):
            kwargs["template"] = json.dumps(prompt_template.template)
        else:
            kwargs["template"] = json.dumps({"<<<genai_template>>>": prompt_template.template})
        return orm.PromptTemplate(**kwargs)

    def _format_dataset_column(
        self,
        dataset: Union[datasets.Dataset, datasets.DatasetDict],
        template_input: str,
        column_name: str,
        batch_size: int = 10,
    ):
        feature_substitutions = set(re.findall(r"\{\{([a-zA-Z0-9_\- ]+)\}\}", template_input))

        def apply_prompts(examples: Dict[str, List]):
            vals = []
            batch_length = min(batch_size, len(examples[list(examples.keys())[0]]))
            for example_idx in range(batch_length):
                args = {
                    feature_name: examples[feature_name][example_idx]
                    for feature_name in examples
                    if feature_name in feature_substitutions
                }
                formatted_template = str(template_input)
                for key, value in args.items():
                    formatted_template = formatted_template.replace("{{" + key + "}}", str(value))
                vals.append(formatted_template)
            examples[column_name] = vals
            return examples

        return dataset.map(apply_prompts, batched=True, batch_size=batch_size)

    def _format_dataset(
        self,
        dataset: Union[datasets.Dataset, datasets.DatasetDict],
        formatting: FormattingTokens,
        new_column_name: str = "genai_input_text",
        expected_output_column_name: Optional[str] = None,
        batch_size: int = 10,
    ):
        if expected_output_column_name:
            # If expected output_column_name is provided, put system and user content in
            # new_column_name, and assistant content in expected_output_column_name.
            template_input = formatting._wrap(
                template=self.template, in_section=True, out_section=False
            )
            dataset = self._format_dataset_column(
                dataset, template_input, new_column_name, batch_size
            )

            template_output = formatting._wrap(
                template=self.template,
                in_section=False,
                out_section=True,
                append_special_tokens=False,
            )
            dataset = self._format_dataset_column(
                dataset, template_output, expected_output_column_name, batch_size
            )
        else:
            # If expected output_column_name is not provided, put all content in
            # new_column_name
            template_input = formatting._wrap(
                template=self.template, in_section=True, out_section=True
            )
            dataset = self._format_dataset_column(
                dataset, template_input, new_column_name, batch_size
            )
        return dataset

    def format(
        self,
        input: Optional[Union[Dict, datasets.Dataset, datasets.DatasetDict]] = None,
        model: Optional[Model] = None,
        formatting_tokens: Optional[FormattingTokens] = None,
        **kwargs: Any,
    ) -> Union[str, datasets.Dataset, datasets.DatasetDict]:
        if formatting_tokens is not None and model is not None:
            raise ValueError("Exactly one of model or formatting can be provided.")
        elif formatting_tokens is None and model is None:
            logging.warning("Using default input formatting.")
            formatting = DefaultFormatting()  # type: FormattingTokens
        else:
            formatting = model.formatting_tokens if model else formatting_tokens

        template = formatting._wrap(self.template)
        if input is None:
            return template
        elif isinstance(input, datasets.Dataset) or isinstance(input, datasets.DatasetDict):
            return self._format_dataset(input, formatting, **kwargs)
        else:
            for key, value in input.items():
                template = template.replace("{{" + key + "}}", str(value))
            return template

Ancestors

pydantic.main.BaseModel

Class variables

var creation_time : Optional[datetime.datetime]
var id : Optional[int]
var model_computed_fields
var model_config
var model_fields
var modified_time : Optional[datetime.datetime]
var name : Optional[str]
var registered_in_genai : bool
var start_prompt_template : bool
var system_generated : bool
var template : Union[str, Dict[str, str]]
var workspace_id : Optional[int]

Methods

def format(self, input: Union[Dict, datasets.arrow_dataset.Dataset, datasets.dataset_dict.DatasetDict, ForwardRef(None)] = None, model: Optional[Model] = None, formatting_tokens: Optional[FormattingTokens] = None, **kwargs: Any) ‑> Union[str, datasets.arrow_dataset.Dataset, datasets.dataset_dict.DatasetDict]

Expand source code

def format(
    self,
    input: Optional[Union[Dict, datasets.Dataset, datasets.DatasetDict]] = None,
    model: Optional[Model] = None,
    formatting_tokens: Optional[FormattingTokens] = None,
    **kwargs: Any,
) -> Union[str, datasets.Dataset, datasets.DatasetDict]:
    if formatting_tokens is not None and model is not None:
        raise ValueError("Exactly one of model or formatting can be provided.")
    elif formatting_tokens is None and model is None:
        logging.warning("Using default input formatting.")
        formatting = DefaultFormatting()  # type: FormattingTokens
    else:
        formatting = model.formatting_tokens if model else formatting_tokens

    template = formatting._wrap(self.template)
    if input is None:
        return template
    elif isinstance(input, datasets.Dataset) or isinstance(input, datasets.DatasetDict):
        return self._format_dataset(input, formatting, **kwargs)
    else:
        for key, value in input.items():
            template = template.replace("{{" + key + "}}", str(value))
        return template

class RegisterHFModelRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class RegisterHFModelRequest(BaseModel):
    hf_model_name: str
    branch: Optional[str] = None
    hf_model_commit: Optional[str] = None
    formatting_tokens: Optional[FormattingTokens] = None
    problem_type: str = ""
    model_config: Optional[Dict[str, Any]] = None

Ancestors

pydantic.main.BaseModel

Class variables

var branch : Optional[str]
var formatting_tokens : Optional[FormattingTokens]
var hf_model_commit : Optional[str]
var hf_model_name : str
var model_computed_fields
var model_config : Optional[Dict[str, Any]]
var model_fields
var problem_type : str

class ResourcePool (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ResourcePool(BaseModel):
    name: str
    gpu_type: Optional[GPUType] = None
    slots_per_agent: int
    slots_available: int
    slots_used: int
    max_agents: int

    @classmethod
    def from_bindings(
        cls,
        bindings: api.bindings.v1ResourcePool,
        extra_metadata: Dict[str, ResourcePoolMetadata] = {},
    ) -> "ResourcePool":
        name = bindings.name
        if name in extra_metadata:
            gpu_type = extra_metadata[name].gpu_type
            max_agents = extra_metadata[name].max_agents
        else:
            gpu_type = GPUType.from_string(bindings.accelerator)
            max_agents = bindings.maxAgents
        slots_per_agent = bindings.slotsPerAgent
        slots_available = bindings.slotsAvailable
        slots_used = bindings.slotsUsed
        return ResourcePool(
            name=name,
            gpu_type=gpu_type,
            slots_per_agent=slots_per_agent,
            slots_available=slots_available,
            slots_used=slots_used,
            max_agents=max_agents,
        )

Ancestors

pydantic.main.BaseModel

Class variables

var gpu_type : Optional[GPUType]
var max_agents : int
var model_computed_fields
var model_config
var model_fields
var name : str
var slots_available : int
var slots_per_agent : int
var slots_used : int

Static methods

def from_bindings(bindings: determined.common.api.bindings.v1ResourcePool, extra_metadata: Dict[str, ResourcePoolMetadata] = {}) ‑> ResourcePool

Expand source code

@classmethod
def from_bindings(
    cls,
    bindings: api.bindings.v1ResourcePool,
    extra_metadata: Dict[str, ResourcePoolMetadata] = {},
) -> "ResourcePool":
    name = bindings.name
    if name in extra_metadata:
        gpu_type = extra_metadata[name].gpu_type
        max_agents = extra_metadata[name].max_agents
    else:
        gpu_type = GPUType.from_string(bindings.accelerator)
        max_agents = bindings.maxAgents
    slots_per_agent = bindings.slotsPerAgent
    slots_available = bindings.slotsAvailable
    slots_used = bindings.slotsUsed
    return ResourcePool(
        name=name,
        gpu_type=gpu_type,
        slots_per_agent=slots_per_agent,
        slots_available=slots_available,
        slots_used=slots_used,
        max_agents=max_agents,
    )

class ResourcePoolInferenceConfigs (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ResourcePoolInferenceConfigs(BaseModel):
    resource_pool: ResourcePool
    vllm_configs: List[VLLMInferenceConfig]

Ancestors

pydantic.main.BaseModel

Class variables

var model_computed_fields
var model_config
var model_fields
var resource_pool : ResourcePool
var vllm_configs : List[VLLMInferenceConfig]

class ResourcePoolMetadata (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ResourcePoolMetadata(BaseModel):
    max_agents: int
    gpu_type: GPUType

Ancestors

pydantic.main.BaseModel

Class variables

var gpu_type : GPUType
var max_agents : int
var model_computed_fields
var model_config
var model_fields

class ResourcePoolTrainingConfigs (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class ResourcePoolTrainingConfigs(BaseModel):
    resource_pool: ResourcePool
    training_configs: List[TrainingConfigRecommendation]

Ancestors

pydantic.main.BaseModel

Class variables

var model_computed_fields
var model_config
var model_fields
var resource_pool : ResourcePool
var training_configs : List[TrainingConfigRecommendation]

class RestartControllerRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class RestartControllerRequest(BaseModel):
    controller_type: ControllerType
    workspace_id: int

    # Currently user's ip (can be located in experiment hyperparameters celery_queue).
    # If not specified, the requester's unique id (ip address) will be used.
    # TODO: Change to username later.
    unique_id: Optional[str] = None
    client_key: Optional[str] = None

Ancestors

pydantic.main.BaseModel

Class variables

var client_key : Optional[str]
var controller_type : ControllerType
var model_computed_fields
var model_config
var model_fields
var unique_id : Optional[str]
var workspace_id : int

class RestartControllerResponse (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class RestartControllerResponse(BaseModel):
    success: bool = None
    prev_experiment_id: Optional[int] = None
    experiment_id: Optional[int] = None

Ancestors

pydantic.main.BaseModel

Class variables

var experiment_id : Optional[int]
var model_computed_fields
var model_config
var model_fields
var prev_experiment_id : Optional[int]
var success : bool

class SampleDatasetRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class SampleDatasetRequest(BaseModel):
    dataset: Optional[Dataset] = None
    dataset_id: Optional[int] = None
    start_index: Optional[int] = None
    number_of_samples: Optional[int] = None
    ratio: Optional[float] = None
    splits: Optional[List[str]] = None
    seed: Optional[int] = None
    as_dataset: bool = False
    workspace_id: int

Ancestors

pydantic.main.BaseModel

Class variables

var as_dataset : bool
var dataset : Optional[Dataset]
var dataset_id : Optional[int]
var model_computed_fields
var model_config
var model_fields
var number_of_samples : Optional[int]
var ratio : Optional[float]
var seed : Optional[int]
var splits : Optional[List[str]]
var start_index : Optional[int]
var workspace_id : int

class StartEndTokens (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class StartEndTokens(BaseModel):
    start: str = ""
    end: str = ""

Ancestors

pydantic.main.BaseModel

Class variables

var end : str
var model_computed_fields
var model_config
var model_fields
var start : str

class TrainingCheckpointUserData (**data: Any)

Custom data stored in checkpoint directories.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class TrainingCheckpointUserData(BaseModel):
    """
    Custom data stored in checkpoint directories.
    """

    base_model_name: str
    """
    The Huggingface repo on which this model is based.  Should not be a local path.
    """

    base_model_revision: str = "main"
    """
    The specific model version (can be a branch name, tag name or commit id).
    """

    base_model_uuid: Optional[str] = None
    """
    The checkpoint UUID this model was based on, if any.  Should not be the UUID of the checkpoint
    itself, but instead of its parent model.
    """

    peft_config: Optional[PEFTConfig] = None
    """
    The PEFT configuration used to train this model, if any.
    """

Ancestors

pydantic.main.BaseModel

Class variables

var base_model_name : str: The Huggingface repo on which this model is based. Should not be a local path.
var base_model_revision : str: The specific model version (can be a branch name, tag name or commit id).
var base_model_uuid : Optional[str]: The checkpoint UUID this model was based on, if any. Should not be the UUID of the checkpoint itself, but instead of its parent model.
var model_computed_fields
var model_config
var model_fields
var peft_config : Optional[PEFTConfig]: The PEFT configuration used to train this model, if any.

class TrainingConfig (*args, **data)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class TrainingConfig(BaseModel):
    data_args: DataArguments = Field(default_factory=DataArguments)
    slots_per_trial: int = 1
    max_steps: int = -1
    num_train_epochs: Optional[float] = None
    save_total_limit: int = 1
    learning_rate: float = 5e-5
    per_device_train_batch_size: int = const.DEFAULT_PER_DEVICE_TRAIN_BATCH_SIZE
    per_device_eval_batch_size: int = const.DEFAULT_PER_DEVICE_EVAL_BATCH_SIZE
    do_train: bool = True
    do_eval: bool = True
    fp16: bool = False  # Deepspeed flag
    bf16: bool = False  # Deepspeed flag
    gradient_checkpointing: bool = const.DEFAULT_TRAIN_GRADIENT_CHECKPOINTING_FLAG
    deepspeed: bool = const.DEFAULT_TRAIN_DEEPSPEED_FLAG
    ds_cpu_offload: bool = const.DEFAULT_TRAIN_DS_CPU_OFF_LOAD
    resource_pool: Optional[str] = None
    seed: int = 1337
    output_dir: str = "/tmp/lore_output"
    eval_steps: Optional[int] = Field(default=None)
    logging_steps: Optional[int] = Field(default=None)
    save_steps: Optional[int] = Field(default=None)
    logging_strategy: str = Field(default="")
    evaluation_strategy: str = Field(default="")
    save_strategy: str = Field(default="")
    save_on_each_node: bool = True

    def __init__(self, *args, **data):
        super().__init__(*args, **data)
        if self.logging_strategy == "":
            self.logging_strategy = "steps"  # The HF default.

        if self.max_steps != -1 and self.num_train_epochs is not None:
            raise ValueError("Only one of `max_steps` or `num_train_epochs` can be specified.")

        if self.max_steps != -1:
            self.evaluation_strategy = (
                self.evaluation_strategy if self.evaluation_strategy != "" else "steps"
            )
            self.eval_steps = (
                self.eval_steps
                if self.eval_steps
                else max(self.max_steps // self.save_total_limit, 1)
            )

            self.save_strategy = self.save_strategy if self.save_strategy != "" else "steps"
            self.save_steps = (
                self.save_steps
                if self.save_steps
                else max(self.max_steps // self.save_total_limit, 1)
            )
            self.logging_steps = (
                self.logging_steps
                if self.logging_steps
                else max(min(500, self.max_steps // self.save_total_limit), 1)
            )
        else:
            self.num_train_epochs = self.num_train_epochs if self.num_train_epochs else 3.0

            self.evaluation_strategy = (
                self.evaluation_strategy if self.evaluation_strategy != "" else "epoch"
            )
            self.eval_steps = (
                self.eval_steps
                if self.eval_steps
                else max(self.num_train_epochs // self.save_total_limit, 1)
            )

            self.save_strategy = self.save_strategy if self.save_strategy != "" else "epoch"
            self.save_steps = (
                self.save_steps
                if self.save_steps
                else max(self.num_train_epochs // self.save_total_limit, 1)
            )

            self.logging_steps = (
                self.logging_steps if self.logging_steps else 500
            )  # The HF default.

Ancestors

pydantic.main.BaseModel

Class variables

var bf16 : bool
var data_args : DataArguments
var deepspeed : bool
var do_eval : bool
var do_train : bool
var ds_cpu_offload : bool
var eval_steps : Optional[int]
var evaluation_strategy : str
var fp16 : bool
var gradient_checkpointing : bool
var learning_rate : float
var logging_steps : Optional[int]
var logging_strategy : str
var max_steps : int
var model_computed_fields
var model_config
var model_fields
var num_train_epochs : Optional[float]
var output_dir : str
var per_device_eval_batch_size : int
var per_device_train_batch_size : int
var resource_pool : Optional[str]
var save_on_each_node : bool
var save_steps : Optional[int]
var save_strategy : str
var save_total_limit : int
var seed : int
var slots_per_trial : int

class TrainingConfigRecommendation (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class TrainingConfigRecommendation(BaseModel):
    slots_per_trial: int
    context_window: int
    batch_size: int
    deepspeed: bool
    gradient_checkpointing: bool
    torch_dtype: str

Ancestors

pydantic.main.BaseModel

Class variables

var batch_size : int
var context_window : int
var deepspeed : bool
var gradient_checkpointing : bool
var model_computed_fields
var model_config
var model_fields
var slots_per_trial : int
var torch_dtype : str

class TrainingConfigRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class TrainingConfigRequest(BaseModel):
    model_id: Optional[int]
    model_architecture: Optional[str]
    # Max number of resource pool to return
    max_resource_pool: Optional[int] = None
    # Max number of config to return per resource pool
    max_config_per_resource_pool: Optional[int] = None

Ancestors

pydantic.main.BaseModel

Class variables

var max_config_per_resource_pool : Optional[int]
var max_resource_pool : Optional[int]
var model_architecture : Optional[str]
var model_computed_fields
var model_config
var model_fields
var model_id : Optional[int]

class UpdatePlaygroundSnaphotRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class UpdatePlaygroundSnaphotRequest(BaseModel):
    # Snapshot id will be from URI, other fields are all optional as user can choose to update
    # any number of them
    id: Optional[int] = None
    name: Optional[str] = None
    model_id: Optional[int] = None
    starting_prompt_template_id: Optional[int] = None
    prompt: Optional[Dict[str, Any]] = None
    dataset_samples: Optional[Dict[str, Any]] = None
    input_dataset_id: Optional[int] = None
    output_dataset_id: Optional[int] = None
    output_datasets: Optional[List[Dataset]] = None
    input_text: Optional[str] = None
    output_text: Optional[str] = None
    generation_config: Optional[Dict[str, Any]] = None
    web: Optional[Dict[str, Any]] = None

Ancestors

pydantic.main.BaseModel

Class variables

var dataset_samples : Optional[Dict[str, Any]]
var generation_config : Optional[Dict[str, Any]]
var id : Optional[int]
var input_dataset_id : Optional[int]
var input_text : Optional[str]
var model_computed_fields
var model_config
var model_fields
var model_id : Optional[int]
var name : Optional[str]
var output_dataset_id : Optional[int]
var output_datasets : Optional[List[Dataset]]
var output_text : Optional[str]
var prompt : Optional[Dict[str, Any]]
var starting_prompt_template_id : Optional[int]
var web : Optional[Dict[str, Any]]

class UploadModelToHFRequest (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class UploadModelToHFRequest(BaseModel):
    model: Model
    hf_repo_owner: str
    hf_repo_name: str
    hf_token: str
    private: bool

Ancestors

pydantic.main.BaseModel

Class variables

var hf_repo_name : str
var hf_repo_owner : str
var hf_token : str
var model : Model
var model_computed_fields
var model_config
var model_fields
var private : bool

class UploadModelToHFResponse (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class UploadModelToHFResponse(BaseModel):
    experiment_id: int

Ancestors

pydantic.main.BaseModel

Class variables

var experiment_id : int
var model_computed_fields
var model_config
var model_fields

class UploadedFiles (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class UploadedFiles(BaseModel):
    file_paths: List[str] = Field(default=[])

Ancestors

pydantic.main.BaseModel

Class variables

var file_paths : List[str]
var model_computed_fields
var model_config
var model_fields

class User (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class User(BaseModel):
    username: str
    admin: bool
    displayName: Optional[str] = None
    remote: Optional[bool] = None

Ancestors

pydantic.main.BaseModel

Class variables

var admin : bool
var displayName : Optional[str]
var model_computed_fields
var model_config
var model_fields
var remote : Optional[bool]
var username : str

class VLLMInferenceConfig (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class VLLMInferenceConfig(BaseModel):
    slots_per_trial: int
    max_new_tokens: int
    batch_size: int
    swap_space: int
    torch_dtype: str

Ancestors

pydantic.main.BaseModel

Class variables

var batch_size : int
var max_new_tokens : int
var model_computed_fields
var model_config
var model_fields
var slots_per_trial : int
var swap_space : int
var torch_dtype : str

class Workspace (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class Workspace(BaseModel):
    id: int
    name: str
    experiment_project_id: int
    controller_project_id: int

    @staticmethod
    def _from_orm(workspace: orm.Workspace) -> "Workspace":
        return Workspace(
            id=workspace.id,
            name=workspace.name,
            experiment_project_id=workspace.experiment_project_id,
            controller_project_id=workspace.controller_project_id,
        )

    async def _to_orm(self) -> orm.Workspace:
        orm_workspace: orm.Workspace = None
        if self.id:
            orm_workspace = await orm.Workspace.objects().get(orm.Workspace.id == self.id)  # type: ignore[arg-type]
        if orm_workspace is None:
            orm_workspace = orm.Workspace(id=self.id)
        orm_workspace.name = self.name
        orm_workspace.experiment_project_id = self.experiment_project_id
        orm_workspace.controller_project_id = self.controller_project_id
        return orm_workspace

Ancestors

pydantic.main.BaseModel

Class variables

var controller_project_id : int
var experiment_project_id : int
var id : int
var model_computed_fields
var model_config
var model_fields
var name : str

class runtimeStreamError (**data: Any)

Usage docs: https://docs.pydantic.dev/2.6/concepts/models/

A base class for creating Pydantic models.

Attributes

__class_vars__: The names of classvars defined on the model.
__private_attributes__: Metadata about the private attributes of the model.
__signature__: The signature for instantiating the model.
__pydantic_complete__: Whether model building is completed, or if there are still undefined fields.
__pydantic_core_schema__: The pydantic-core schema used to build the SchemaValidator and SchemaSerializer.
__pydantic_custom_init__: Whether the model has a custom __init__ function.
__pydantic_decorators__: Metadata containing the decorators defined on the model. This replaces Model.__validators__ and Model.__root_validators__ from Pydantic V1.
__pydantic_generic_metadata__: Metadata for generic models; contains data used for a similar purpose to args, origin, parameters in typing-module generics. May eventually be replaced by these.
__pydantic_parent_namespace__: Parent namespace of the model, used for automatic rebuilding of models.
__pydantic_post_init__: The name of the post-init method for the model, if defined.
__pydantic_root_model__: Whether the model is a RootModel.
__pydantic_serializer__: The pydantic-core SchemaSerializer used to dump instances of the model.
__pydantic_validator__: The pydantic-core SchemaValidator used to validate instances of the model.
__pydantic_extra__: An instance attribute with the values of extra fields from validation when model_config['extra'] == 'allow'.
__pydantic_fields_set__: An instance attribute with the names of fields explicitly set.
__pydantic_private__: Instance attribute with the values of private attributes set on the model instance.

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Expand source code

class runtimeStreamError(BaseModel):
    message: str = Field(default="")

Ancestors

pydantic.main.BaseModel

Class variables

var message : str
var model_computed_fields
var model_config
var model_fields