### Installation and Environment Setup

Source: https://context7.com/aiola-lab/drax/llms.txt

Instructions for setting up the Drax environment, including conda environment creation, PyTorch installation with CUDA, and Drax installation.

```APIDOC
## Installation and Environment Setup

Instructions for setting up the Drax environment.

### Conda Environment Setup

```bash
# Create conda environment
conda create -n drax python=3.10 -y
conda activate drax
```

### PyTorch Installation (with CUDA)

```bash
pip install -U pip
pip install torch==2.7.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
```

### Drax Installation

```bash
# Install Drax
pip install -e .

# For systems without pre-installed torch
pip install -e .[with-torch]
```

### Verification

```bash
# Verify installation
python -c "from drax import Transcriber; print('Drax installed successfully')"
```
```

--------------------------------

### Drax Installation and Environment Setup

Source: https://context7.com/aiola-lab/drax/llms.txt

Instructions for setting up the development environment for Drax. This includes creating a conda environment, installing PyTorch with CUDA support, and installing the Drax library itself.

```bash
# Create conda environment
conda create -n drax python=3.10 -y
conda activate drax

# Install PyTorch with CUDA support
pip install -U pip
pip install torch==2.7.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128

# Install Drax
pip install -e .

# For systems without pre-installed torch
pip install -e .[with-torch]

# Verify installation
python -c "from drax import Transcriber; print('Drax installed successfully')"
```

--------------------------------

### Install Drax and Dependencies

Source: https://github.com/aiola-lab/drax/blob/main/README.md

Installs the Drax library and its core dependencies, including PyTorch and torchaudio. It also provides an option to install with PyTorch pre-installed for convenience.

```bash
conda create -n drax python=3.10 -y
conda activate drax
pip install -U pip
pip install torch==2.7.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
pip install -e .

pip install -e .[with-torch]
```

--------------------------------

### Setup Pre-commit Hooks with Ruff

Source: https://github.com/aiola-lab/drax/blob/main/README.md

Installs pre-commit hooks and Ruff for code linting and formatting. This is part of the development setup to maintain code quality.

```bash
pip install pre-commit ruff
pre-commit install
# Run manually
pre-commit run -a
```

--------------------------------

### Source Distribution API

Source: https://context7.com/aiola-lab/drax/llms.txt

API for controlling the initial noise distribution in discrete flow matching, with examples for uniform and masked distributions.

```APIDOC
## Source Distribution API

Controls the initial noise distribution for discrete flow matching.

### Uniform Source Distribution

```python
from drax.flow.source_dist import UniformSourceDistribution
import torch

# Uniform distribution (random tokens)
uniform_dist = UniformSourceDistribution(vocab_size=51865)
x0_uniform = uniform_dist.sample((4, 128), device="cuda")  # [batch, seq_len]
# Each token randomly initialized from [0, vocab_size)
```

### Masked Source Distribution

```python
from drax.flow.source_dist import MaskedSourceDistribution
import torch

# Masked distribution (all tokens start as mask token)
mask_token_id = 51864
masked_dist = MaskedSourceDistribution(mask_token=mask_token_id)
x0_masked = masked_dist.sample((4, 128), device="cuda")
# All tokens initialized to mask_token_id
```

### Usage with Transcriber

```python
# Use with Transcriber
asr_masked = Transcriber(
    model_path="aiola-labs/drax-v1",
    source_distribution=masked_dist,
    device="cuda"
)
```

### Sampling like existing tensor

```python
# Sample like existing tensor
reference = torch.randint(0, 100, (8, 64))
new_sample = uniform_dist.sample_like(reference)  # matches shape and device
```
```

--------------------------------

### Command-Line Transcription Script

Source: https://context7.com/aiola-lab/drax/llms.txt

This section details how to use the `generate.py` script for command-line audio transcription, including basic and advanced usage examples.

```APIDOC
## Command-Line Transcription Script

Standalone script for quick transcription from the command line.

### Usage Examples

```bash
# Basic transcription
python generate.py \
    --model_path aiola-labs/drax-v1 \
    --audio_path /path/to/audio.wav \
    --language en \
    --sampling_steps 16 \
    --temperature 0.01

# Advanced usage with longer sequences
python generate.py \
    --model_path ./local/model/checkpoint \
    --audio_path recording.mp3 \
    --language de \
    --sampling_steps 32 \
    --temperature 0.05 \
    --seq_length 256
```

### Programmatic Usage

```python
from generate import main
import argparse

args = argparse.Namespace(
    model_path="aiola-labs/drax-v1",
    audio_path="/data/interview.wav",
    language="en",
    temperature=1e-2,
    sampling_steps=24,
    seq_length=128
)

result = main(args)
print(f"Transcript: {result[0].transcript}")
print(f"Audio path: {result[0].audio_path}")
print(f"Language: {result[0].language}")
```
```

--------------------------------

### Source Distribution API in Drax

Source: https://context7.com/aiola-lab/drax/llms.txt

Controls the initial noise distribution for discrete flow matching. This Python code demonstrates initializing uniform and masked source distributions and using them with a Transcriber.

```python
from drax.flow.source_dist import UniformSourceDistribution, MaskedSourceDistribution
import torch

# Uniform distribution (random tokens)
uniform_dist = UniformSourceDistribution(vocab_size=51865)
x0_uniform = uniform_dist.sample((4, 128), device="cuda")  # [batch, seq_len]
# Each token randomly initialized from [0, vocab_size)

# Masked distribution (all tokens start as mask token)
mask_token_id = 51864
MaskedSourceDistribution = MaskedSourceDistribution(mask_token=mask_token_id)
x0_masked = masked_dist.sample((4, 128), device="cuda")
# All tokens initialized to mask_token_id

# Use with Transcriber
asr_masked = Transcriber(
    model_path="aiola-labs/drax-v1",
    source_distribution=masked_dist,
    device="cuda"
)

# Sample like existing tensor
reference = torch.randint(0, 100, (8, 64))
new_sample = uniform_dist.sample_like(reference)  # matches shape and device
```

--------------------------------

### Programmatic Transcription with Drax

Source: https://context7.com/aiola-lab/drax/llms.txt

Programmatic usage of the Drax transcription script. This snippet demonstrates how to import and use the `main` function from the `generate` module with predefined arguments and prints the resulting transcript.

```python
# Programmatic usage of the script
from generate import main
import argparse

args = argparse.Namespace(
    model_path="aiola-labs/drax-v1",
    audio_path="/data/interview.wav",
    language="en",
    temperature=1e-2,
    sampling_steps=24,
    seq_length=128
)

result = main(args)
print(f"Transcript: {result[0].transcript}")
print(f"Audio path: {result[0].audio_path}")
print(f"Language: {result[0].language}")
```

--------------------------------

### Command-Line Transcription with Drax

Source: https://context7.com/aiola-lab/drax/llms.txt

Standalone script for quick transcription from the command line. It supports basic and advanced usage with configurable model paths, audio files, language, sampling steps, and temperature.

```bash
# Basic transcription
python generate.py \
    --model_path aiola-labs/drax-v1 \
    --audio_path /path/to/audio.wav \
    --language en \
    --sampling_steps 16 \
    --temperature 0.01

# Advanced usage with longer sequences
python generate.py \
    --model_path ./local/model/checkpoint \
    --audio_path recording.mp3 \
    --language de \
    --sampling_steps 32 \
    --temperature 0.05 \
    --seq_length 256
```

--------------------------------

### Model Saving and Loading API

Source: https://context7.com/aiola-lab/drax/llms.txt

API for saving and loading model checkpoints, including support for safetensors and loading from HuggingFace Hub.

```APIDOC
## Model Saving and Loading

Serialization API for model checkpoints with safetensors support.

### Saving a Model

```python
from drax.model.drax_model import Drax
import torch

# Save model
model = Drax.from_pretrained("aiola-labs/drax-v1")
model.save_pretrained("/path/to/checkpoint")
# Creates:
#   /path/to/checkpoint/config.json
#   /path/to/checkpoint/encoder.safetensors
#   /path/to/checkpoint/decoder.safetensors
```

### Loading a Model

```python
# Load from local directory
model_local = Drax.from_pretrained(
    "/path/to/checkpoint",
    device=torch.device("cuda"),
    torch_dtype=torch.float16
)

# Load from HuggingFace Hub with caching
model_hub = Drax.from_pretrained(
    "aiola-labs/drax-v1",
    cache_dir="/custom/cache/dir",
    device="cuda"
)
```

### Accessing Model Components

```python
# Access model components
print(f"Decoder vocab size: {model_hub.decoder.vocab_size}")
print(f"Hidden size: {model_hub.decoder.config.hidden_size}")
print(f"Number of blocks: {model_hub.decoder.config.n_blocks}")
print(f"Supported languages: {model_hub.decoder.config.support_language_codes}")
```
```

--------------------------------

### Model Saving and Loading in Drax

Source: https://context7.com/aiola-lab/drax/llms.txt

Serialization API for model checkpoints with safetensors support. This Python code shows how to save a model to a local directory, load it from the local directory, and load from the HuggingFace Hub with custom caching.

```python
from drax.model.drax_model import Drax
import torch

# Save model
model = Drax.from_pretrained("aiola-labs/drax-v1")
model.save_pretrained("/path/to/checkpoint")
# Creates:
#   /path/to/checkpoint/config.json
#   /path/to/checkpoint/encoder.safetensors
#   /path/to/checkpoint/decoder.safetensors

# Load from local directory
model_local = Drax.from_pretrained(
    "/path/to/checkpoint",
    device=torch.device("cuda"),
    torch_dtype=torch.float16
)

# Load from HuggingFace Hub with caching
model_hub = Drax.from_pretrained(
    "aiola-labs/drax-v1",
    cache_dir="/custom/cache/dir",
    device="cuda"
)

# Access model components
print(f"Decoder vocab size: {model_hub.decoder.vocab_size}")
print(f"Hidden size: {model_hub.decoder.config.hidden_size}")
print(f"Number of blocks: {model_hub.decoder.config.n_blocks}")
print(f"Supported languages: {model_hub.decoder.config.support_language_codes}")
```

--------------------------------

### Transformer Decoder Configuration

Source: https://context7.com/aiola-lab/drax/llms.txt

Details on the low-level transformer decoder architecture, including configuration options and forward pass usage with audio cross-attention.

```APIDOC
## Transformer Decoder Configuration

Low-level decoder architecture with audio cross-attention.

### Configuration and Initialization

```python
from drax.model.transformer import Transformer
from omegaconf import OmegaConf
import torch

# Define decoder configuration
config = OmegaConf.create({
    "vocab_size": 51865,
    "hidden_size": 1024,
    "n_heads": 16,
    "n_blocks": 24,
    "cond_dim": 1024,
    "whisper_dim": 1280,  # Whisper encoder output dimension
    "dropout": 0.1,
    "length": 128,
    "support_language_codes": ["en", "es", "fr", "de", "it", "pt", "ja", "ko", "zh"]
})

# Initialize transformer decoder
decoder = Transformer(
    vocab_size=config.vocab_size,
    masked=False,  # use uniform source distribution
    config=config
)
decoder.to("cuda").eval()
```

### Forward Pass with Audio Conditioning

```python
# Forward pass with audio conditioning
batch_size, seq_len = 2, 128
x_t = torch.randint(0, 51865, (batch_size, seq_len)).cuda() # Input tokens
time = torch.rand(batch_size).cuda() # Timestep
audio_embeddings = torch.randn(batch_size, 1500, 1280).cuda()  # Whisper encoder output

logits = decoder(
    x_t=x_t,
    time=time,
    audio_embeddings=audio_embeddings,
    preserve_mask=None,
    cfg_strength=1.5,  # classifier-free guidance strength
    audio_drop_prob=0.0  # audio dropout during training
)
# Returns: Tensor[batch, seq_len, vocab_size]
```
```

--------------------------------

### Generate Transcriptions with DraxMixtureDiscreteEulerSolver

Source: https://context7.com/aiola-lab/drax/llms.txt

The DraxMixtureDiscreteEulerSolver API facilitates transcription generation through iterative refinement using discrete flow matching. It requires a Drax model instance, a scheduler, and vocabulary size. The solver can sample initial states, set prompt tokens to preserve, and generate transcription samples given cached audio data and a step size.

```python
import torch
from flow_matching.path.scheduler import PolynomialConvexScheduler
from drax.flow.solver import DraxMixtureDiscreteEulerSolver
from drax.flow.source_dist import UniformSourceDistribution

# Initialize solver components
vocab_size = 51865  # Whisper vocabulary size
scheduler = PolynomialConvexScheduler(n=1.0)  # controls interpolation path
source_dist = UniformSourceDistribution(vocab_size=vocab_size)

solver = DraxMixtureDiscreteEulerSolver(
    model=model,  # Drax model instance
    scheduler=scheduler,
    vocabulary_size=vocab_size
)

# Sample initial state from source distribution
batch_size, seq_length = 2, 128
x_init = source_dist.sample((batch_size, seq_length), device="cuda")

# Set prompt tokens (language, task, etc.) that should be preserved
forced_ids = torch.tensor([50258, 50259, 50359, 50363])  # <|startoftranscript|><|en|><|transcribe|><|notimestamps|>
x_init[:, :len(forced_ids)] = forced_ids
preserve_mask = torch.zeros((batch_size, seq_length), dtype=torch.bool, device="cuda")
preserve_mask[:, :len(forced_ids)] = True

# Generate transcription samples
samples = solver.sample(
    x_init=x_init,
    step_size=1/32,  # 32 sampling steps
    audio_projected=audio_cache["audio_projected"],
    audio_k_all=audio_cache["audio_k_all"],
    audio_v_all=audio_cache["audio_v_all"],
    preserve_mask=preserve_mask,
    temperature=1e-2,
    verbose=True  # show progress bar
)
# Returns: Tensor[batch, seq_length] of token IDs
```

--------------------------------

### Batch Audio Transcription with Drax

Source: https://github.com/aiola-lab/drax/blob/main/README.md

Processes multiple audio files simultaneously for transcription. Supports specifying a list of audio paths and corresponding languages for batch inference.

```python
audio_paths = ["/path/to/audio1.wav", "/path/to/audio2.wav"]
languages = ["en", "de"]
result = asr.transcribe(audio_paths, language=languages)
print(result.transcript)
```

--------------------------------

### Transformer Decoder Configuration in Drax

Source: https://context7.com/aiola-lab/drax/llms.txt

Low-level decoder architecture with audio cross-attention. This Python code defines a decoder configuration using OmegaConf and initializes a Transformer decoder, followed by a forward pass with audio conditioning.

```python
from drax.model.transformer import Transformer
from omegaconf import OmegaConf
import torch

# Define decoder configuration
config = OmegaConf.create({
    "vocab_size": 51865,
    "hidden_size": 1024,
    "n_heads": 16,
    "n_blocks": 24,
    "cond_dim": 1024,
    "whisper_dim": 1280,  # Whisper encoder output dimension
    "dropout": 0.1,
    "length": 128,
    "support_language_codes": ["en", "es", "fr", "de", "it", "pt", "ja", "ko", "zh"]
})

# Initialize transformer decoder
decoder = Transformer(
    vocab_size=config.vocab_size,
    masked=False,  # use uniform source distribution
    config=config
)
decoder.to("cuda").eval()

# Forward pass with audio conditioning
batch_size, seq_len = 2, 128
x_t = torch.randint(0, 51865, (batch_size, seq_len)).cuda()
time = torch.rand(batch_size).cuda()
audio_embeddings = torch.randn(batch_size, 1500, 1280).cuda()  # Whisper encoder output

logits = decoder(
    x_t=x_t,
    time=time,
    audio_embeddings=audio_embeddings,
    preserve_mask=None,
    cfg_strength=1.5,  # classifier-free guidance strength
    audio_drop_prob=0.0  # audio dropout during training
)
# Returns: Tensor[batch, seq_len, vocab_size]
```

--------------------------------

### Advanced Audio Transcription with Drax Options

Source: https://github.com/aiola-lab/drax/blob/main/README.md

Transcribes audio with customizable sampling steps and temperature for controlling the generation process. This allows for fine-tuning the transcription output.

```python
result = asr.transcribe("/path/to/audio.wav", language="en", sampling_steps=32, temperature=1e-2)
print(result[0].transcript)
```

--------------------------------

### DraxMixtureDiscreteEulerSolver API

Source: https://context7.com/aiola-lab/drax/llms.txt

Implements the discrete flow matching solver for generating transcriptions through iterative refinement of token sequences.

```APIDOC
## DraxMixtureDiscreteEulerSolver API

### Description
Discrete flow matching solver that generates transcriptions through iterative refinement.

### Method
`DraxMixtureDiscreteEulerSolver` class

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None (Initialization parameters)
- **model** (Drax) - Required - An instance of the Drax model.
- **scheduler** (PolynomialConvexScheduler) - Required - Controls the interpolation path between the source and target distributions.
- **vocabulary_size** (int) - Required - The size of the vocabulary (e.g., 51865 for Whisper).

### Methods
- **`sample(x_init, step_size, audio_projected, audio_k_all, audio_v_all, preserve_mask, temperature, verbose)`**: Generates transcription samples by iteratively refining an initial state.
  - **x_init** (torch.Tensor) - Initial token sequence sampled from the source distribution. Shape: `[batch, seq_length]`.
  - **step_size** (float) - The step size for the Euler solver (e.g., 1/32 for 32 sampling steps).
  - **audio_projected** (torch.Tensor) - Cached projected audio embeddings from the Drax model.
  - **audio_k_all** (torch.Tensor) - Cached audio attention keys from the Drax model.
  - **audio_v_all** (torch.Tensor) - Cached audio attention values from the Drax model.
  - **preserve_mask** (torch.Tensor) - A boolean mask indicating which initial tokens should be preserved throughout the sampling process. Shape: `[batch, seq_length]`.
  - **temperature** (float) - Controls the randomness during sampling.
  - **verbose** (bool) - If True, displays a progress bar during sampling.
  - Returns: `torch.Tensor` - The final sampled token IDs. Shape: `[batch, seq_length]`.

### Request Example
```python
import torch
from flow_matching.path.scheduler import PolynomialConvexScheduler
from drax.flow.solver import DraxMixtureDiscreteEulerSolver
from drax.flow.source_dist import UniformSourceDistribution

# Assuming 'model' and 'audio_cache' are already defined from Drax Model API

vocab_size = 51865
scheduler = PolynomialConvexScheduler(n=1.0)
source_dist = UniformSourceDistribution(vocab_size=vocab_size)

solver = DraxMixtureDiscreteEulerSolver(
    model=model, 
    scheduler=scheduler,
    vocabulary_size=vocab_size
)

batch_size, seq_length = 2, 128
x_init = source_dist.sample((batch_size, seq_length), device="cuda")

# Example prompt tokens (e.g., language, task)
forced_ids = torch.tensor([50258, 50259, 50359, 50363]) # <|startoftranscript|><|en|><|transcribe|><|notimestamps|>
x_init[:, :len(forced_ids)] = forced_ids
preserve_mask = torch.zeros((batch_size, seq_length), dtype=torch.bool, device="cuda")
preserve_mask[:, :len(forced_ids)] = True

samples = solver.sample(
    x_init=x_init,
    step_size=1/32,
    audio_projected=audio_cache["audio_projected"],
    audio_k_all=audio_cache["audio_k_all"],
    audio_v_all=audio_cache["audio_v_all"],
    preserve_mask=preserve_mask,
    temperature=1e-2,
    verbose=True
)
```

### Response
#### Success Response (200)
- **samples** (torch.Tensor) - A tensor containing the generated token IDs for the transcription. Shape: `[batch, seq_length]`.
```

--------------------------------

### Basic Audio Transcription with Drax

Source: https://github.com/aiola-lab/drax/blob/main/README.md

Performs speech-to-text transcription on a single audio file using the Drax Transcriber. Requires specifying the model path and the language of the audio.

```python
from drax import Transcriber

asr = Transcriber(model_path="aiola-labs/drax-v1")  # HF repo or local path
result = asr.transcribe("/path/to/audio.wav", language="en")
print(result[0].transcript)
```

--------------------------------

### Transcriber API

Source: https://context7.com/aiola-lab/drax/llms.txt

Provides the main interface for transcribing audio files, supporting both single file and batch processing with multilingual capabilities.

```APIDOC
## Transcriber API

### Description
Main interface for transcribing audio files with support for single and batch processing.

### Method
`transcribe`

### Parameters
#### Path Parameters
None

#### Query Parameters
- **audio_path** (str or list[str]) - Required - Path to the audio file(s) to transcribe.
- **language** (str or list[str]) - Optional - The language(s) of the audio file(s). Defaults to auto-detection.
- **temperature** (float) - Optional - Controls the randomness of the transcription. Lower values lead to more deterministic output.
- **sampling_steps** (int) - Optional - The number of discrete flow matching steps to perform. Typical values range from 16 to 64.
- **seq_length** (int) - Optional - The maximum sequence length for the decoder.

### Request Example
```python
from drax import Transcriber

asr = Transcriber(model_path="aiola-labs/drax-v1")

# Single file transcription
result = asr.transcribe(audio_path="/path/to/audio.wav", language="en", temperature=1e-2, sampling_steps=32)
print(result[0].transcript)

# Batch processing
audio_paths = ["/path/to/audio1.wav", "/path/to/audio2.wav"]
languages = ["en", "de"]
results = asr.transcribe(audio_paths, language=languages, sampling_steps=16)

for sample in results:
    print(f"Transcript: {sample.transcript}")
```

### Response
#### Success Response (200)
A list of `TranscriptionResult` objects, where each object contains:
- **audio_path** (str) - The path to the processed audio file.
- **language** (str) - The detected or specified language of the audio.
- **transcript** (str) - The transcribed text.

#### Response Example
```json
[
  {
    "audio_path": "/path/to/audio.wav",
    "language": "en",
    "transcript": "Hello, this is a test recording."
  }
]
```
```

--------------------------------

### Transcribe Audio Files with Drax Transcriber API

Source: https://context7.com/aiola-lab/drax/llms.txt

The Transcriber API provides the main interface for transcribing single audio files or batches of files. It allows specifying the audio path, language, and sampling parameters like temperature and number of steps. The output is a list of transcription results, each containing the audio path, detected language, and the transcribed text.

```python
from drax import Transcriber

# Initialize transcriber with a pretrained model
asr = Transcriber(
    model_path="aiola-labs/drax-v1",  # HuggingFace repo or local path
    device="cuda"  # auto-detected if not specified
)

# Single audio file transcription
result = asr.transcribe(
    audio_path="/path/to/audio.wav",
    language="en",
    temperature=1e-2,  # controls randomness (lower = more deterministic)
    sampling_steps=32,  # number of flow matching steps (16-64 typical)
    seq_length=128     # max sequence length for decoder
)
print(result[0].transcript)  # "Hello, this is a test recording."

# Batch processing with multiple languages
audio_paths = [
    "/path/to/english.wav",
    "/path/to/german.wav",
    "/path/to/spanish.wav"
]
languages = ["en", "de", "es"]
results = asr.transcribe(audio_paths, language=languages, sampling_steps=16)

# Iterate through results
for sample in results:
    print(f"File: {sample.audio_path}")
    print(f"Language: {sample.language}")
    print(f"Transcript: {sample.transcript}\n")
```

--------------------------------

### Drax Model API for End-to-End ASR

Source: https://context7.com/aiola-lab/drax/llms.txt

The Drax Model API combines an audio encoder and a flow decoder for end-to-end ASR. It allows loading pretrained models, processing audio features to obtain embeddings, and building an audio cache for efficient inference. The model can then be used for forward passes with cached audio data to generate logits representing token probabilities.

```python
import torch
from drax.model.drax_model import Drax

# Load pretrained model
model = Drax.from_pretrained(
    "aiola-labs/drax-v1",
    device=torch.device("cuda"),
    torch_dtype=torch.bfloat16
)
model.eval()

# Process audio features (expects Whisper-preprocessed input)
audio_features = torch.randn(1, 80, 3000)  # [batch, mel_bins, time]
audio_embeddings = model.encode_audio(audio_features)  # [batch, seq_len, hidden]

# Build reusable audio cache for efficient inference
audio_cache = model.build_audio_cache(audio_embeddings)
# Returns: {
#   "audio_projected": Tensor[B, L, H],
#   "audio_k_all": Tensor[B, num_blocks, L, H],
#   "audio_v_all": Tensor[B, num_blocks, L, H]
# }

# Forward pass with cached audio
x_t = torch.randint(0, 51865, (1, 128))  # current token state
t = torch.tensor([0.5])  # time step in [0, 1]
logits = model(
    x=x_t,
    t=t,
    audio_projected=audio_cache["audio_projected"],
    audio_k_all=audio_cache["audio_k_all"],
    audio_v_all=audio_cache["audio_v_all"],
    temperature=1e-1
)
# Returns softmax probabilities over vocabulary
```

--------------------------------

### Drax Model API

Source: https://context7.com/aiola-lab/drax/llms.txt

The core model class that integrates an audio encoder (Whisper) and a discrete flow matching decoder for end-to-end speech recognition.

```APIDOC
## Drax Model API

### Description
Core model class that combines audio encoder and flow decoder for end-to-end ASR.

### Method
`Drax` class

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None (Initialization parameters)
- **model_path** (str) - Required - HuggingFace repository ID or local path to the pretrained model.
- **device** (torch.device) - Optional - The device to load the model on (e.g., 'cuda', 'cpu'). Auto-detected if not specified.
- **torch_dtype** (torch.dtype) - Optional - The data type for model weights (e.g., `torch.bfloat16`).

### Methods
- **`from_pretrained(model_path, device, torch_dtype)`**: Class method to load a pretrained Drax model.
- **`encode_audio(audio_features)`**: Encodes audio features into embeddings.
  - **audio_features** (torch.Tensor) - Input audio features (e.g., Mel spectrograms). Expected shape: `[batch, mel_bins, time]`.
  - Returns: `torch.Tensor` - Audio embeddings. Shape: `[batch, seq_len, hidden]`.
- **`build_audio_cache(audio_embeddings)`**: Builds a reusable cache from audio embeddings for efficient inference.
  - **audio_embeddings** (torch.Tensor) - Output from `encode_audio`.
  - Returns: `dict` - A dictionary containing cached audio projections and attention keys/values.
    - **`audio_projected`** (Tensor[B, L, H])
    - **`audio_k_all`** (Tensor[B, num_blocks, L, H])
    - **`audio_v_all`** (Tensor[B, num_blocks, L, H])

### Forward Pass
- **`model(x, t, audio_projected, audio_k_all, audio_v_all, temperature)`**: Performs a forward pass through the decoder.
  - **x** (torch.Tensor) - Current token state. Shape: `[batch, seq_length]`.
  - **t** (torch.Tensor) - Time step in the flow matching process. Shape: `[batch]` (values in [0, 1]).
  - **audio_projected** (torch.Tensor) - Cached projected audio embeddings.
  - **audio_k_all** (torch.Tensor) - Cached audio attention keys.
  - **audio_v_all** (torch.Tensor) - Cached audio attention values.
  - **temperature** (float) - Controls the randomness of the output probabilities.
  - Returns: `torch.Tensor` - Logits (softmax probabilities) over the vocabulary. Shape: `[batch, seq_length, vocab_size]`.

### Request Example
```python
import torch
from drax.model.drax_model import Drax

model = Drax.from_pretrained("aiola-labs/drax-v1", device=torch.device("cuda"))
model.eval()

audio_features = torch.randn(1, 80, 3000)
audio_embeddings = model.encode_audio(audio_features)
audio_cache = model.build_audio_cache(audio_embeddings)

x_t = torch.randint(0, 51865, (1, 128))
t = torch.tensor([0.5])
logits = model(
    x=x_t,
    t=t,
    audio_projected=audio_cache["audio_projected"],
    audio_k_all=audio_cache["audio_k_all"],
    audio_v_all=audio_cache["audio_v_all"],
    temperature=1e-1
)
```

### Response
#### Success Response (200)
- **logits** (torch.Tensor) - Softmax probabilities over the vocabulary for the next token prediction. Shape: `[batch, seq_length, vocab_size]`.
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.