Clef Omni 30B A3B

"""Clef Omni 30B A3B, generated from a declarative template."""

from __future__ import annotations

import base64
import io
import sys
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator

from runware_serverless import serve

MODEL = "Cloudflare/clef-omni"
REVISION = "f32dfc060386af3eadec676c64b34f24ece2bcde"
CACHE = "/weights"


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 | None = Field(
        default=None,
        description="The question to answer about the state, in words. Optional, where the criteria speak for themselves.",
    )
    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. Not read for noul, which is always false against true.",
    )

    @model_validator(mode="after")
    def _criteria_fit_the_type(self) -> "Question":
        # Checked here rather than left to the model, which raises on the same
        # case once the request is already being encoded.
        if self.type != "noul" and not self.criteria:
            raise ValueError("a %s question needs criteria: the options it scores" % self.type)

        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.",
    )
    images: list[str] | None = Field(
        default=None,
        description="Images to read alongside the state, each a base64 image file.",
    )
    audio: list[str] | None = Field(
        default=None,
        description="Recordings to listen to alongside the state, each a base64 audio file.",
    )


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.",
    )
    usage: dict[str, Any] = Field(description="How many tokens the state and the questions came to.")


@serve
class Decision:
    def load(self) -> None:
        if hasattr(self, "model"):
            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

        local = snapshot_download(MODEL, revision=REVISION, cache_dir=CACHE)

        # The joint schema head is not a transformers class, and the module that
        # assembles it with the backbone ships inside the checkpoint. Putting
        # the snapshot on the path and importing from it is what the model card
        # documents, and there is no other published way in.
        if local not in sys.path:
            sys.path.insert(0, local)

        from joint_schema_model import load_release_model, systemone

        self.systemone = systemone
        # Straight onto the card: the loader passes device_map for the backbone
        # and moves the head itself. A worker's host RAM is half its card's
        # VRAM, so a model assembled on the host and moved has a ceiling below
        # the card's own.
        #
        # This release loads the thinker alone. The talker and the vocoder are
        # in the files and are left there, which is why the card holds less than
        # the download suggests.
        self.model, self.processor = load_release_model(local, device="cuda")

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

        if request.images:
            body["images"] = [self._image(encoded) for encoded in request.images]
        if request.audio:
            # Bytes rather than a decoded waveform: this release's reader takes
            # encoded audio and resamples it to the rate the model wants, so
            # handing it an array would mean getting that rate right here.
            body["audio"] = [base64.b64decode(clip) for clip in request.audio]

        return Decided(**{
            key: value
            for key, value in self.systemone(self.model, self.processor, body).items()
            if key in ("answers", "usage")
        })

    @staticmethod
    def _image(encoded: str) -> Any:
        from PIL import Image

        # Decoded to pixels here rather than passed along as base64. The
        # processor takes an image, a path or a URL, and a URL would be this
        # worker reaching out mid-request.
        image = Image.open(io.BytesIO(base64.b64decode(encoded)))

        return image.convert("RGB")

configVersion: 1
id: clef-omni-30b-a3b
name: Clef Omni 30B A3B
source:
  type: code
  entry: model.py
compute:
  gpuType: rtx-pro-6000
build:
  requirements:
    - torch==2.14.1
    - torchvision==0.29.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - pillow==13.0.0
    - huggingface-hub==2.4.0
    - av==16.0.1
volumes:
  - /weights