Voxtral Mini 3B

"""Voxtral Mini 3B, generated from a declarative template."""

from __future__ import annotations

import base64
import tempfile

from pydantic import BaseModel, ConfigDict, Field

from runware_serverless import serve

MODEL = "mistralai/Voxtral-Mini-3B-2507"
REVISION = "3060fe34b35ba5d44202ce9ff3c097642914f8f3"
CACHE = "/weights"
SAMPLE_RATE = 16000
DEFAULT_LANGUAGE = None


class Request(BaseModel):
    # An input 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")

    audio: str = Field(description="The speech to transcribe, as a base64 audio file. Any container soundfile reads, so WAV, FLAC and OGG all work.")
    language: str | None = Field(
        default=DEFAULT_LANGUAGE,
        description="The language to transcribe as, as a two-letter code. Left out, the model detects it.",
    )
    max_tokens: int = Field(default=512, ge=1, le=4096, description="Upper bound on the transcript's length in tokens.")


class Transcribed(BaseModel):
    text: str = Field(description="The transcript.")
    seconds: float = Field(description="How long the audio ran, after decoding.")
    language: str | None = Field(description="The language it was transcribed as, or null where the model chose.")


@serve
class Transcription:
    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.
        import torch
        from transformers import VoxtralForConditionalGeneration, VoxtralProcessor

        self.torch = torch
        self.processor = VoxtralProcessor.from_pretrained(
            MODEL, revision=REVISION, cache_dir=CACHE
        )
        # Straight onto the card as it loads. 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.
        self.model = VoxtralForConditionalGeneration.from_pretrained(
            MODEL,
            revision=REVISION,
            cache_dir=CACHE,
            dtype=torch.bfloat16,
            device_map="cuda",
        )

    def predict(self, request: Request) -> Transcribed:
        import soundfile

        raw = base64.b64decode(request.audio)

        # Written to a file and passed as a path, not decoded into an array.
        # Given a path the processor calls load_audio_as itself, with
        # force_mono and the sampling rate it wants, so it does the conversion
        # the model needs. Given an array it requires a "format" argument and
        # refuses without one, and any resampling is then ours to get right.
        with tempfile.NamedTemporaryFile(suffix=".wav") as handle:
            handle.write(raw)
            handle.flush()

            info = soundfile.info(handle.name)

            inputs = self.processor.apply_transcription_request(
                audio=handle.name,
                model_id=MODEL,
                language=request.language,
                sampling_rate=SAMPLE_RATE,
            ).to("cuda", dtype=self.torch.bfloat16)

            generated = self.model.generate(**inputs, max_new_tokens=request.max_tokens)

            # Only the tokens the model added. The request template is long,
            # and decoding the whole sequence returns the prompt with the
            # transcript buried at the end of it.
            new = generated[:, inputs.input_ids.shape[1]:]
            text = self.processor.batch_decode(new, skip_special_tokens=True)[0]

        return Transcribed(
            text=text.strip(),
            seconds=info.frames / info.samplerate,
            language=request.language,
        )

configVersion: 1
id: voxtral-mini-3b
name: Voxtral Mini 3B
source:
  type: code
  entry: model.py
compute:
  gpuType: l40s
build:
  requirements:
    - torch==2.14.1
    - transformers==5.21.0
    - accelerate==1.15.0
    - soundfile==0.14.0
    - librosa==0.11.0
    - mistral-common==1.13.1
volumes:
  - /weights