Laya

"""Laya, generated from a declarative template."""

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from runware_serverless import serve

MODEL = "convaiinnovations/laya"
REVISION = "7b928d828b7b0e022f929d9bd2e44165aa270148"
CACHE = "/weights"
CHECKPOINTS = ["english","multilingual","typed-decisions"]


class Question(BaseModel):
    # A key this model does not read is a line the caller believed was doing
    # something, so it is refused rather than dropped.
    model_config = ConfigDict(extra="forbid")

    type: Literal["choice", "score", "noul"] = Field(
        description="choice picks one of the labels in criteria, score returns a number on the scale criteria describes, and noul returns the probability that the answer is yes.",
    )
    instructions: str = Field(
        min_length=1,
        description="The question to answer about the state, in words.",
    )
    criteria: dict[str, str] | list[str] | None = Field(
        default=None,
        description="For choice, the allowed labels: either a mapping of label to what it means, or a plain list of labels. For score, the scale as an ordered list of levels, lowest first. Optional for noul, where it may describe what false and true mean.",
    )

    @model_validator(mode="after")
    def _criteria_fit_the_type(self) -> "Question":
        # Checked here rather than left to the model, which raises on the same
        # cases one request later and from three frames down.
        if self.type == "choice" and not self.criteria:
            raise ValueError("a choice question needs criteria: the labels it may answer with")
        if self.type == "score" and not isinstance(self.criteria, list):
            raise ValueError("a score question takes criteria as an ordered list of levels, lowest first")

        return self


class Request(BaseModel):
    # As above: a misspelled field is refused rather than silently ignored.
    model_config = ConfigDict(extra="forbid")

    state: str | dict[str, Any] | list[Any] = Field(
        description="What to decide about: text, a JSON object such as a ticket or an invoice, or a list of conversation turns.",
    )
    questions: dict[str, Question] = Field(
        min_length=1,
        description="The questions to answer, keyed by an id of your own. Every id comes back in the answers.",
    )
    checkpoint: Literal["auto", "english", "multilingual", "typed-decisions"] = Field(
        default="auto",
        description="Which checkpoint answers. Leave it on auto to have the script and language of the state decide, which is what the English and multilingual pair is for. The response reports which one answered.",
    )
    max_len: int | None = Field(
        default=None,
        ge=1,
        le=8192,
        description="How many tokens of the state to read. Only the multilingual checkpoint reads past 1024, up to 8192, and the response reports whether the state fit.",
    )
    min_confidence: float | None = Field(
        default=None,
        ge=0.0,
        le=1.0,
        description="Below this calibrated confidence an answer is reported as undecided rather than guessed.",
    )


class Decided(BaseModel):
    answers: dict[str, dict[str, Any]] = Field(
        description="One answer per question id. A choice answer carries choice, confidence and probabilities, a score answer score, confidence, legend and probabilities, and a noul answer the probability that it is true.",
    )
    checkpoint: str | None = Field(description="Which checkpoint answered.")
    usage: dict[str, Any] = Field(
        description="Token counts, and whether the state fit: truncated, state_tokens and state_tokens_dropped.",
    )


@serve
class Decision:
    def load(self) -> None:
        if hasattr(self, "router"):
            return

        # Imported here rather than at the top of the file. The build imports
        # this file with the platform's own dependencies, so nothing above
        # exists until a worker runs.
        from huggingface_hub import snapshot_download
        from laya import Router

        # Snapshotted here, and the Router given local paths, because the
        # package has no cache_dir anywhere in its surface: left to itself it
        # downloads to huggingface_hub's default, outside the volume, and every
        # cold start fetches the weights again.
        #
        # The pattern list is the one the package uses for a single checkpoint,
        # plus the two bundled subfolders. Each checkpoint carries its own
        # encoder config, so nothing is fetched from the backbone's repository
        # at load time; what the list leaves out is the model card's images.
        local = snapshot_download(
            MODEL,
            revision=REVISION,
            cache_dir=CACHE,
            allow_patterns=[
                "rl_agent_config.json",
                "model.safetensors",
                "tokenizer/*",
                "encoder/*",
                "multilingual/*",
                "typed-decisions/*",
            ],
        )

        self.router = Router(
            # A local path and a subfolder, in place of the repository id the
            # package would otherwise resolve itself.
            models={
                "english": (local, None),
                "multilingual": (local, "multilingual"),
                "typed-decisions": (local, "typed-decisions"),
            },
            device="cuda",
            # All three resident. Routing is microseconds and a cold build is
            # seconds, and the default cap of two would evict a checkpoint on
            # any request that names the third.
            preload=True,
            max_loaded=len(CHECKPOINTS),
            # Off, rather than left at its default of "auto". vegaml is a core
            # dependency of the package, so "auto" is on, and that lane fetches
            # a separate repository the first time a state is longer than the
            # encoders can read: a model download in the middle of a request.
            # Off, a long state is read as far as max_len reaches and the
            # response reports the rest as dropped.
            vega=False,
        )

    def predict(self, request: Request) -> Decided:
        # exclude_none so a question that named no criteria is absent rather
        # than present and null, which is what the model checks for.
        questions = {
            name: question.model_dump(exclude_none=True)
            for name, question in request.questions.items()
        }

        extra: dict[str, Any] = {}
        if request.checkpoint != "auto":
            extra["model"] = request.checkpoint
        if request.max_len is not None:
            extra["max_len"] = request.max_len
        if request.min_confidence is not None:
            extra["min_confidence"] = request.min_confidence

        result = self.router.predict(request.state, questions, **extra)

        return Decided(
            answers=result.get("answers") or {},
            checkpoint=(result.get("routing") or {}).get("model"),
            usage=result.get("usage") or {},
        )

configVersion: 1
id: laya-decisions
name: Laya
source:
  type: code
  entry: model.py
compute:
  gpuType: l40s
build:
  requirements:
    - laya==0.4.29
    - torch==2.14.1
    - transformers==5.21.0
    - huggingface-hub==2.4.0
volumes:
  - /weights