EasyDeL EasyData

https://easydel.readthedocs.io/en/latest/
Admin
EasyData is a comprehensive data management module providing flexible tools for loading,...
Tokens:107,913
Snippets:1,127
Trust Score:9.6
Update:6 months ago
Show doc for...
Tokens:
Raw
### Run Whisper Server using Example Script

Source: https://easydel.readthedocs.io/en/latest/whisper_api.html

Execute the provided example script to start the Whisper server. This is an alternative to using the CLI or direct programmatic execution.

```bash
python examples/whisper_server_example.py --model "openai/whisper-large-v3-turbo" --port 8000

```

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

### Run Whisper Server using Example Script

Source: https://easydel.readthedocs.io/en/latest/_sources/whisper_api.md

Execute the example script to start the Whisper server. Requires specifying the model and port.

```bash
python examples/whisper_server_example.py --model "openai/whisper-large-v3-turbo" --port 8000
```

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

### Perform Large-Scale Pre-training Setup

Source: https://easydel.readthedocs.io/en/latest/easydata/README.html

Example of loading and mixing datasets for large-scale pre-training workflows.

```python
from datasets import load_dataset
from easydel.data import (
    block_mixture_interleave,
    MixedShardedSource,
    HFDatasetShardedSource,
    WeightScheduler,
    WeightSchedulePoint,
)
import easydel as ed

# Load datasets with streaming
code_ds = load_dataset("bigcode/starcoderdata", split="train", streaming=True)
text_ds = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)

# Simple mixing
mixed = block_mixture_interleave(
    [code_ds, text_ds],
    weights={"code": 0.3, "text": 0.7},
    block_size=1000,
    seed=42,
    stop="restart",
)
```

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

### Initialize and Train Reward Model with RewardTrainer

Source: https://easydel.readthedocs.io/en/latest/trainers/reward.html

This example shows how to load a tokenizer and model, define RewardConfig, initialize RewardTrainer, and start the training process. Ensure the tokenizer has a pad_token_id and the model is configured for sequence classification with one label.

```python
import easydel as ed
import jax
from jax import numpy as jnp
from transformers import AutoTokenizer
from datasets import load_dataset

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load dataset with preference pairs
dataset = load_dataset(
    "trl-lib/ultrafeedback_binarized",
    split="train[:5%]"  # Using a small subset for demonstration
)

# Load model
model = ed.AutoEasyDeLModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    num_labels=1,  # Reward model outputs a single score
    dtype=jnp.bfloat16
)

# Create Reward config
config = ed.RewardConfig(
    model_name="reward_model_example",
    save_directory="reward_checkpoints",
    max_sequence_length=2048,
    center_rewards_coefficient=0.1,
    total_batch_size=16,
    learning_rate=1e-6,
    learning_rate_end=5e-7,
    num_train_epochs=3,
    use_wandb=True,
)

# Initialize trainer
trainer = ed.RewardTrainer(
    arguments=config,
    model=model,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

# Start training
trainer.train()

```

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

### Use Example Client for Transcription

Source: https://easydel.readthedocs.io/en/latest/_sources/whisper_api.md

Run the example client script to transcribe an audio file. Specify the server URL, mode, language, and whether to include timestamps.

```bash
python examples/whisper_client_example.py audio.mp3 --server http://localhost:8000 --mode transcribe --language en --timestamps
```

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

### Perform Basic KTO Training

Source: https://easydel.readthedocs.io/en/latest/easydata/trainer_integration.html

Initializes a KTO trainer setup.

```python
from datasets import load_dataset
import easydel as ed
```

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

### Install EasyDeL on Kaggle or Colab

Source: https://easydel.readthedocs.io/en/latest/_sources/install.md

Commands to prepare the environment by removing conflicting packages and installing EasyDeL with JAX and PyTorch dependencies.

```shell
pip uninstall torch-xla -y -q  # Remove pre-installed torch-xla (for TPUs)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu -qU  # Install PyTorch for model conversion
pip install git+https://github.com/erfanzar/easydel -qU  # Install EasyDeL from the latest source
pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -qU  # Install JAX for TPUs
```

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

### Install Dependencies

Source: https://easydel.readthedocs.io/en/latest/_sources/trc-welcome.md

Installs PyTorch and the EasyDeL library using eopod.

```shell
# Install required dependencies
eopod run pip install torch --index-url https://download.pytorch.org/whl/cpu

# Install EasyDeL from the latest source
eopod run pip install git+https://github.com/erfanzar/easydel
```

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

### Install and Configure EOpod

Source: https://easydel.readthedocs.io/en/latest/_sources/install.md

Setup eopod for managing TPU hosts and ensure the local binary path is correctly exported to the shell profile.

```shell
pip install eopod
```

```shell
echo 'export PATH="$HOME/.local/bin:$PATH"' >> ~/.bashrc
source ~/.bashrc  # Apply changes immediately
```

```shell
eopod configure --project-id YOUR_PROJECT_ID --zone YOUR_ZONE --tpu-name YOUR_TPU_NAME
```

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

### Install EasyDeL with Cloud Storage Support

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/README.md

Install the easydel package. For cloud storage capabilities, install with the appropriate extras, e.g., easydel[gcs] for Google Cloud Storage or easydel[s3] for Amazon S3.

```bash
pip install easydel

# For cloud storage
pip install easydel[gcs]  # Google Cloud Storage
pip install easydel[s3]   # Amazon S3
```

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

### Install FastAPI and EasyDeL Dependencies

Source: https://easydel.readthedocs.io/en/latest/whisper_api.html

Install necessary libraries for the Whisper API server. This includes FastAPI with all extras, uvicorn for running the server, and the easydel library.

```bash
pip install "fastapi[all]" uvicorn easydel

```

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

### Install EasyDeL on TPU Pods

Source: https://easydel.readthedocs.io/en/latest/_sources/index.rst

Use this script for distributed installation on TPU pods. Specify your TPU type and choose between PyPI or GitHub source.

```bash
python -m easydel.scripts.install_on_hosts --tpu-type v4-16 --source github
```

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

### Initialize and Run RewardTrainer

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/reward.md

Full workflow for loading a model, configuring the trainer, and starting the training process.

```python
import easydel as ed
import jax
from jax import numpy as jnp
from transformers import AutoTokenizer
from datasets import load_dataset

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load dataset with preference pairs
dataset = load_dataset(
    "trl-lib/ultrafeedback_binarized",
    split="train[:5%]"  # Using a small subset for demonstration
)

# Load model
model = ed.AutoEasyDeLModelForSequenceClassification.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    num_labels=1,  # Reward model outputs a single score
    dtype=jnp.bfloat16
)

# Create Reward config
config = ed.RewardConfig(
    model_name="reward_model_example",
    save_directory="reward_checkpoints",
    max_sequence_length=2048,
    center_rewards_coefficient=0.1,
    total_batch_size=16,
    learning_rate=1e-6,
    learning_rate_end=5e-7,
    num_train_epochs=3,
    use_wandb=True,
)

# Initialize trainer
trainer = ed.RewardTrainer(
    arguments=config,
    model=model,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

# Start training
trainer.train()
```

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

### Start Easydel API Server

Source: https://easydel.readthedocs.io/en/latest/esurge_examples.html

Initialize and launch the Easydel API server. API keys can be optionally configured for security.

```python
# Initialize engine
engine = ed.eSurge(
    model="meta-llama/Llama-3.2-3B-Instruct",
    max_model_len=2048,
    max_num_seqs=16,
    hbm_utilization=0.85,
)

# Enable monitoring
engine.start_monitoring()

# Launch API server (keys optional)
server = ed.eSurgeApiServer(
    engine,
    require_api_key=True,
    api_keys={"demo-key": {"label": "dashboard"}},
)
server.fire(host="0.0.0.0", port=8000)
```

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

### Use Example Client for Transcription

Source: https://easydel.readthedocs.io/en/latest/whisper_api.html

Utilize the example client script to transcribe an audio file via the Whisper API server. This command specifies the audio file, server address, mode, and language.

```bash
python examples/whisper_client_example.py audio.mp3 --server http://localhost:8000 --mode transcribe --language en --timestamps

```

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

### Large-Scale Pre-training Example

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/README.md

This example demonstrates loading datasets from HuggingFace Hub with streaming, mixing them using block_mixture_interleave with specified weights and block size, and setting a stop condition.

```python
from datasets import load_dataset
from easydel.data import (
    block_mixture_interleave,
    MixedShardedSource,
    HFDatasetShardedSource,
    WeightScheduler,
    WeightSchedulePoint,
)
import easydel as ed

# Load datasets with streaming
code_ds = load_dataset("bigcode/starcoderdata", split="train", streaming=True)
text_ds = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)

# Simple mixing
mixed = block_mixture_interleave(
    [code_ds, text_ds],
    weights={"code": 0.3, "text": 0.7},
    block_size=1000,
    seed=42,
    stop="restart",
)

```

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

### Initialize and Use SFTTrainer

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/sft.md

A complete example demonstrating loading a tokenizer, dataset, and model, configuring SFTTrainer, and initiating the training process. Ensure the tokenizer has a pad_token_id set.

```python
import easydel as ed
import jax
from jax import numpy as jnp
from transformers import AutoTokenizer
from datasets import load_dataset

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load dataset
dataset = load_dataset(
    "Anthropic/hh-rlhf",
    data_dir="helpful-base",
    split="train[:10%]"  # Using a small subset for demonstration
)

# Load model
model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype=jnp.bfloat16
)

# Create SFT config
config = ed.SFTConfig(
    model_name="sft_example",
    save_directory="sft_checkpoints",
    dataset_text_field="chosen",  # Field containing the text to train on
    packing=True,                 # Enable sequence packing for efficiency
    learning_rate=2e-5,
    learning_rate_end=5e-6,
    total_batch_size=16,
    num_train_epochs=3,
    use_wandb=True,
    num_of_sequences=512,         # For packing
)

# Initialize trainer
trainer = ed.SFTTrainer(
    arguments=config,
    model=model,
    train_dataset=dataset,
    processing_class=tokenizer,
)

# Start training
trainer.train()
```

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

### Install eSurge with TPU support

Source: https://easydel.readthedocs.io/en/latest/_sources/esurge.rst

Install the easydel library with TPU specific optimizations. This command installs the necessary packages for leveraging TPUs with eSurge.

```bash
pip install easydel[tpu]
```

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

### SFT Training Pipeline

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/pipeline.md

Example of a full SFT training pipeline configuration and execution loop.

```python
from easydel.data import (
    Pipeline,
    PipelineConfig,
    DatasetConfig,
    PackStageConfig,
    LoadStageConfig,
)

config = PipelineConfig(
    datasets=[
        DatasetConfig(
            data_files="conversations/*.jsonl",
            tokenizer="meta-llama/Llama-2-7b-chat-hf",
            content_field="messages",  # Chat data
        ),
    ],
    pack=PackStageConfig(
        enabled=True,
        seq_length=2048,
        eos_token_id=2,
    ),
    load=LoadStageConfig(
        batch_size=8,
        shuffle_buffer_size=10000,
    ),
)

pipeline = Pipeline.from_config(config)
for batch in pipeline.source().tokenize().pack().load().build():
    sft_train_step(batch)
```

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

### Configure a Full Pipeline with Stage Settings

Source: https://easydel.readthedocs.io/en/latest/easydata/pipeline.html

Illustrates the comprehensive PipelineConfig object, including dataset definitions and global/stage-specific settings. Use this for detailed pipeline setup.

```python
from easydel.data import (
    PipelineConfig,
    DatasetConfig,
    TokenizeStageConfig,
    MixStageConfig,
    PackStageConfig,
    LoadStageConfig,
    SaveStageConfig,
)

config = PipelineConfig(
    # Required: at least one dataset
    datasets=[
        DatasetConfig(
            name="dataset1",
            data_files="data/*.parquet",
            tokenizer="meta-llama/Llama-2-7b",
        ),
    ],

    # Global settings
    default_tokenizer="meta-llama/Llama-2-7b",  # Fallback tokenizer
    streaming=True,  # Enable streaming mode
    seed=42,  # Random seed

    # Stage configurations
    tokenize=TokenizeStageConfig(max_length=2048),
    mix=MixStageConfig(weights={"ds1": 0.5, "ds2": 0.5}),
    pack=PackStageConfig(enabled=True, seq_length=2048),
    load=LoadStageConfig(batch_size=8),
    save=SaveStageConfig(enabled=True, output_dir="./output"),
)


```

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

### Install Packages on TPU Hosts

Source: https://easydel.readthedocs.io/en/latest/_sources/install.md

Use eopod to distribute package installations across all TPU slices.

```shell
eopod run pip install torch --index-url https://download.pytorch.org/whl/cpu -qU  # Required for model conversion
eopod run pip install git+https://github.com/erfanzar/easydel -qU  # Install EasyDeL from the latest source
eopod run pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -qU  # Install JAX for TPUs
```

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

### Initialize and Start EasyDel API Server

Source: https://easydel.readthedocs.io/en/latest/_sources/esurge_examples.rst

Initializes the eSurge engine with specified model and parameters, enables monitoring, and launches the API server. API keys can be configured for access control.

```python
    # Initialize engine
    engine = ed.eSurge(
        model="meta-llama/Llama-3.2-3B-Instruct",
        max_model_len=2048,
        max_num_seqs=16,
        hbm_utilization=0.85,
    )

    # Enable monitoring
    engine.start_monitoring()

    # Launch API server (keys optional)
    server = ed.eSurgeApiServer(
        engine,
        require_api_key=True,
        api_keys={"demo-key": {"label": "dashboard"}},
    )
    server.fire(host="0.0.0.0", port=8000)
```

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

### Run Whisper Server using CLI

Source: https://easydel.readthedocs.io/en/latest/whisper_api.html

Start the EasyDeL Whisper API server using the command-line interface. Specify the model to use and the port for the server to listen on.

```bash
python -m easydel.inference.vwhisper.cli --model "openai/whisper-large-v3-turbo" --port 8000

```

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

### Configure DPOTrainer with IPO Loss

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/dpo.md

Example of configuring DPOTrainer to use the Implicit Preference Optimization (IPO) loss function. Adjust beta as needed for different loss types.

```python
config = ed.DPOConfig(
    loss_type="ipo",
    beta=0.2,  # May need different beta for different loss types
    # Other parameters...
)
```

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

### Initialize and Use DiskCache

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/caching.md

Set up a persistent disk cache with options for compression and automatic expiry. Supports put, get, and contains operations, with `contains` also checking for expiry.

```python
from easydel.data import DiskCache

cache = DiskCache(
    cache_dir="./disk_cache",
    compression="zstd",      # Best compression ratio
    expiry_seconds=86400,    # Auto-expire after 24 hours
)

# Operations
cache.put("key", large_data)
result = cache.get("key")  # Returns (data, metadata) or None

# Manual expiry check
if cache.contains("key"):  # Also checks expiry
    result = cache.get("key")
```

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

### Run Whisper Server Programmatically

Source: https://easydel.readthedocs.io/en/latest/_sources/whisper_api.md

Start the Whisper server using the EasyDeL module. Requires importing easydel and jax.numpy. Specify model, host, port, and dtype.

```python
import easydel as ed
from jax import numpy as jnp

ed.inference.vwhisper.run_server(
    model_name="openai/whisper-large-v3-turbo",
    host="0.0.0.0",
    port=8000,
    dtype=jnp.bfloat16
)
```

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

### Data Format Examples for Trainers

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/trainer_integration.md

Illustrates the expected data formats for different Easydel trainers, including SFT (text/messages), DPO/ORPO/CPO (chosen/rejected), KTO (paired/unpaired), GRPO (prompts), and Reward (prompt/chosen/rejected).

```python
# SFT: text or messages
{"text": "..."} or {"messages": [...]}

# DPO/ORPO/CPO: chosen/rejected
{"chosen": [...], "rejected": [...]} 

# KTO: paired or unpaired
{"chosen": [...], "rejected": [...]}  # Paired (unpaired internally)
{"prompt": "...", "completion": "...", "label": True/False}  # Already unpaired

# GRPO: prompts
{"prompt": "..."} or {"chosen": [...], "rejected": [...]} 

# Reward: chosen/rejected
{"prompt": "...", "chosen": "...", "rejected": "..."}
```

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

### Initialize ORPOTrainer

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/trainer_integration.md

Sets up the ORPO trainer for preference optimization.

```python
from datasets import load_dataset
import easydel as ed

dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = ed.ORPOTrainer(
    model=model,
    train_dataset=dataset,
    processing_class=tokenizer,
    arguments=ed.ORPOConfig(
        max_prompt_length=512,
        max_completion_length=512,
        beta=0.1,
    ),
)
trainer.train()
```

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

### Install eopod CLI

Source: https://easydel.readthedocs.io/en/latest/_sources/trc-welcome.md

Installs the command-line tool required for managing TPU pods.

```shell
pip install eopod
```

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

### Initialize and Use DPOTrainer

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/dpo.md

Load tokenizer, dataset, model, and reference model to initialize the DPOTrainer. Ensure tokenizer has a pad token ID.

```python
import easydel as ed
from transformers import AutoTokenizer
from datasets import load_dataset
import jax

# Load tokenizer and prepare dataset
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load dataset with preference pairs
dataset = load_dataset(
    "trl-lib/ultrafeedback_binarized",
    split="train[:10%]"  # Using a small subset for demonstration
)

# Load model and reference model
model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype=jax.numpy.bfloat16,
  # ... model loading configs
)
ref_model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype=jax.numpy.bfloat16,
  # ... model loading configs
)

# Create DPO config
config = ed.DPOConfig(
    model_name="dpo_example",
    save_directory="dpo_checkpoints",
    beta=0.1,
    loss_type="sigmoid",
    max_length=2048,
    max_prompt_length=1024,
    total_batch_size=16,
    learning_rate=1e-6,
    num_train_epochs=3,
    use_wandb=True,
)

# Initialize the trainer
trainer = ed.DPOTrainer(
    model=model,
    reference_model=ref_model,
    arguments=config,
    train_dataset=dataset,
    processing_class=tokenizer,
)

# Start training
trainer.train()
```

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

### Install Whisper Server Dependencies

Source: https://easydel.readthedocs.io/en/latest/_sources/whisper_api.md

Install FastAPI and other necessary dependencies for the Whisper server using pip.

```bash
pip install "fastapi[all]" uvicorn easydel
```

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

### Filter Out Examples with Empty Fields

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/transforms.md

FilterNonEmpty removes examples where any of the specified fields are empty strings or None. Useful for ensuring data completeness.

```python
from easydel.data import FilterNonEmpty

transform = FilterNonEmpty(["text", "response"])

result = transform({"text": "hello", "response": "world"})  # Returns example
result = transform({"text": "", "response": "world"})       # Returns None
```

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

### Basic GRPO Trainer Usage for Fine-tuning

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/grpo.md

Set up and use the GRPOTrainer by loading a tokenizer and model, preparing a dataset, defining a reward function, configuring GRPO, and initializing the trainer.

```python
import easydel as ed
import jax
from jax import numpy as jnp
from transformers import AutoTokenizer
from datasets import load_dataset

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load model
model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype=jnp.bfloat16
)

# Load dataset (must have a "prompt" field)
dataset = load_dataset("gsm8k", "main", split="train[:10%]")

# Define a reward function - must return values between 0 and 1
def math_problem_reward_function(completion, reference_answer=None):
    """Example reward function that extracts and verifies an answer."""
    try:
        # Extract answer number from completion
        import re
        answer_match = re.search(r"The answer is (\d+)", completion)
        if answer_match:
            extracted_answer = int(answer_match.group(1))

            # Compare with reference if available
            if reference_answer is not None:
                return 1.0 if extracted_answer == reference_answer else 0.0

            # Some other scoring logic when reference not available
            return min(1.0, max(0.0, 0.5))  # Default score
        return 0.0
    except:
        return 0.0

# Create GRPO config
config = ed.GRPOConfig(
    model_name="grpo_math_solver",
    save_directory="grpo_checkpoints",
    beta=0.04,
    max_prompt_length=512,
    max_completion_length=256,
    total_batch_size=16,
    learning_rate=1e-6,
    num_train_epochs=3,
    use_wandb=True,
)

# Initialize trainer
trainer = ed.GRPOTrainer(
    arguments=config,
    model=model,
    reward_funcs=math_problem_reward_function,
    train_dataset=dataset,
    processing_class=tokenizer,
)

# Start training
trainer.train()
```

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

### Initialize and Use ORPOTrainer in Python

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/orpo.md

Load tokenizer, dataset, and model, then create an ORPOConfig and ORPOTrainer instance. Ensure tokenizer has a pad_token_id and the dataset is loaded correctly. The model should be loaded with the appropriate dtype.

```python
import easydel as ed
import jax
from jax import numpy as jnp
from transformers import AutoTokenizer
from datasets import load_dataset

# Load tokenizer and prepare it
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

# Load dataset with preference pairs
dataset = load_dataset(
    "trl-lib/ultrafeedback_binarized",
    split="train[:5%]"  # Using a small subset for demonstration
)

# Load model
model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct",
    dtype=jnp.bfloat16
)

# Create ORPO config
config = ed.ORPOConfig(
    model_name="orpo_example",
    save_directory="orpo_checkpoints",
    beta=0.1,
    max_length=2048,
    max_prompt_length=1024,
    max_completion_length=1024,
    total_batch_size=16,
    learning_rate=1e-6,
    learning_rate_end=5e-7,
    num_train_epochs=3,
    use_wandb=True,
)

# Initialize trainer
trainer = ed.ORPOTrainer(
    arguments=config,
    model=model,
    train_dataset=dataset,
    processing_class=tokenizer,
)

# Start training
trainer.train()
```

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

### Tracking Data Source within Mixed Examples

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/mixing.md

Access the automatically added '__source__' field in examples from `MixedShardedSource` to identify the origin dataset and apply source-specific processing.

```python
mixed = MixedShardedSource(sources, ...)

for example in mixed.open_shard(mixed.shard_names[0]):
    source_name = example.get("__source__")  # Added automatically
    if source_name == "code":
        # Apply code-specific processing
        pass
```

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

### Initialize and Access Data Sources

Source: https://easydel.readthedocs.io/en/latest/easydata/pipeline.html

Demonstrates initializing a pipeline and accessing the loaded data sources. Use this to verify data loading after the source stage.

```python
pipeline = Pipeline.from_config(config)
pipeline.source()

# Access loaded sources
sources = pipeline.get_data()
# {"dataset1": ParquetShardedSource, "dataset2": JsonShardedSource}


```

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

### Iterate through Parquet shards and examples

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/sources.md

Access data by iterating through shard names and then opening each shard to process individual examples. This method is suitable for large datasets that are split into multiple files.

```python
# Iteration
for shard_name in source.shard_names:
    for example in source.open_shard(shard_name):
        print(example)
```

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

### Complete Chatbot Application Example

Source: https://easydel.readthedocs.io/en/latest/_sources/esurge_examples.rst

A complete example of a chatbot application using EasyDel's eSurge engine. It includes setting environment variables, initializing the engine, formatting conversation history, and processing user input.

```python
    import os
    os.environ["EASYDEL_AUTO"] = "1"

    from jax import numpy as jnp
    import easydel as ed

    class ChatBot:
        def __init__(self, model_id="microsoft/phi-2"):
            self.engine = ed.eSurge(
                model=model_id,
                max_model_len=2048,
                max_num_seqs=8,
                hbm_utilization=0.8,
                esurge_name="chatbot",
            )
            self.conversation = []

        def format_prompt(self, messages):
            """Format conversation for model."""
            prompt = ""
            for msg in messages:
                role = msg["role"].capitalize()
                prompt += f"{role}: {msg['content']}\n"
            prompt += "Assistant: "
            return prompt

        def chat(self, user_input):
            """Process user input and return response."""
            self.conversation.append({"role": "user", "content": user_input})

            prompt = self.format_prompt(self.conversation)


```

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

### Initialize Trainer Preprocessing Transforms

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/trainer_integration.md

Import and instantiate trainer-specific preprocessing transforms like SFTPreprocessTransform and DPOPreprocessTransform. Ensure the tokenizer is provided and configure relevant length parameters.

```python
from easydel.trainers.transforms import (
    SFTPreprocessTransform,
    DPOPreprocessTransform,
    ORPOPreprocessTransform,
    KTOPreprocessTransform,
    GRPOPreprocessTransform,
    RewardPreprocessTransform,
    BCOPreprocessTransform,
    CPOPreprocessTransform,
)

# SFT transform
sft_transform = SFTPreprocessTransform(
    tokenizer=tokenizer,
    max_length=2048,
)

# DPO transform
dpo_transform = DPOPreprocessTransform(
    tokenizer=tokenizer,
    max_prompt_length=512,
    max_completion_length=512,
)
```

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

### Run DPO Training from Command Line

Source: https://easydel.readthedocs.io/en/latest/trainers/dpo.html

Execute DPO training using the easydel.scripts.finetune.dpo module from the command line, specifying model, dataset, and various training parameters.

```bash
python -m easydel.scripts.finetune.dpo \
  --repo_id meta-llama/Llama-3.1-8B-Instruct \
  --dataset_name trl-lib/ultrafeedback_binarized \
  --dataset_split "train[:90%]" \
  --refrence_model_repo_id meta-llama/Llama-3.1-8B-Instruct \
  --attn_mechanism vanilla \
  --beta 0.08 \
  --loss_type sigmoid \
  --max_length 2048 \
  --max_prompt_length 1024 \
  --total_batch_size 16 \
  --learning_rate 1e-6 \
  --log_steps 50 \
  --num_train_epochs 3 \
  --do_last_save \
  --save_steps 1000 \
  --use_wandb

```

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

### Pre-tokenization Pipeline

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/pipeline.md

Example of a pipeline that only tokenizes and saves data to disk.

```python
config = PipelineConfig(
    datasets=[
        DatasetConfig(
            data_files="raw_data/*.jsonl",
            tokenizer="meta-llama/Llama-2-7b",
            save_path="./tokenized_data",
        ),
    ],
    tokenize=TokenizeStageConfig(max_length=2048),
    save=SaveStageConfig(
        enabled=True,
        format="parquet",
        compression="zstd",
    ),
)

# Just tokenize and save
Pipeline.from_config(config).source().tokenize().save().build()
```

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

### Configure TrainingArguments

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/base_trainer.md

Set up training parameters like epochs, learning rate, batch size, and output directory. Use `max_steps` to override `num_train_epochs` and `gradient_checkpointing` for memory optimization.

```python
from easydel.trainers import TrainingArguments

args = TrainingArguments(
    model_name="MyFineTunedModel",
    num_train_epochs=3,
    learning_rate=2e-5,
    total_batch_size=32, # Effective batch size (per_device_train_batch_size * num_devices)
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=1,
    weight_decay=0.01,
    max_steps=-1, # If > 0, overrides num_train_epochs
    gradient_checkpointing="nothing_saveable", # Or "everything_saveable"
    sharding_array=(1, -1, 1, 1, 1), # Device mesh shape for model sharding
    use_fast_kernels=True,
    logging_steps=10,
    save_steps=100, # Save checkpoint every 100 steps
    save_total_limit=2, # Keep only the latest 2 checkpoints
    evaluation_strategy="steps", # Or "epoch"
    eval_steps=100, # Evaluate every 100 steps
    output_dir="runs/my_model_run"
    # Add other relevant TrainingArguments as needed
)

# Example usage (assuming MyTrainerSubclass exists):
# trainer = MyTrainerSubclass(arguments=args, ...)
# trainer.train()
```

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

### JSON Data Format Examples

Source: https://easydel.readthedocs.io/en/latest/easydata/sources.html

Comparison of JSONL and JSON array formats for data loading.

```json
{"text": "First example"}
{"text": "Second example"}
{"text": "Third example"}
```

```json
[
  {"text": "First example"},
  {"text": "Second example"},
  {"text": "Third example"}
]
```

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

### Customize Dataset Preprocessing

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/reward.md

Example of a custom preprocessing function to format preference pairs for the RewardTrainer.

```python
def preprocess_function(examples):
    # Custom preprocessing logic
    chosen_texts = [prompt + chosen for prompt, chosen in zip(examples["prompt"], examples["chosen"])]
    rejected_texts = [prompt + rejected for prompt, rejected in zip(examples["prompt"], examples["rejected"])]

    # Tokenize both chosen and rejected
    chosen_inputs = tokenizer(chosen_texts, truncation=True, max_length=1024)
    rejected_inputs = tokenizer(rejected_texts, truncation=True, max_length=1024)

    return {
        "input_ids_chosen": chosen_inputs["input_ids"],
        "attention_mask_chosen": chosen_inputs["attention_mask"],
        "input_ids_rejected": rejected_inputs["input_ids"],
        "attention_mask_rejected": rejected_inputs["attention_mask"],
    }

# Process the dataset
processed_dataset = dataset.map(
    preprocess_function,
    batched=True,
    num_proc=4,
    remove_columns=dataset.column_names
)

# Use the processed dataset with the trainer
trainer = ed.RewardTrainer(
    arguments=config,
    model=model,
    train_dataset=processed_dataset,
    tokenizer=tokenizer,
)
```

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

### Monitoring Streaming Throughput

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/streaming.md

Track processing speed by calculating examples per second during iteration.

```python
import time
from easydel.data import ParquetShardedSource

source = ParquetShardedSource("gs://bucket/data/*.parquet")

start = time.time()
count = 0
for shard in source.shard_names:
    for example in source.open_shard(shard):
        count += 1
        if count % 10000 == 0:
            elapsed = time.time() - start
            rate = count / elapsed
            print(f"Processed {count} examples, {rate:.1f} ex/sec")
```

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

### Multi-Dataset Pre-training Pipeline

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/pipeline.md

Example of a pre-training pipeline using multiple datasets with weighted mixing.

```python
config = PipelineConfig(
    datasets=[
        DatasetConfig(
            name="code",
            data_files="gs://bucket/code/*.parquet",
            tokenizer="bigcode/starcoder",
        ),
        DatasetConfig(
            name="text",
            data_files="gs://bucket/text/*.parquet",
            tokenizer="meta-llama/Llama-2-7b",
        ),
        DatasetConfig(
            name="math",
            data_files="gs://bucket/math/*.parquet",
            tokenizer="meta-llama/Llama-2-7b",
        ),
    ],
    mix=MixStageConfig(
        weights={"code": 0.4, "text": 0.5, "math": 0.1},
        block_size=2000,
    ),
    pack=PackStageConfig(
        enabled=True,
        seq_length=4096,
        strategy="pool",
    ),
    load=LoadStageConfig(batch_size=32),
)

pipeline = Pipeline.from_config(config)
for batch in pipeline.source().tokenize().mix().pack().load().build():
    pretrain_step(batch)
```

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

### Initialize and Run Chatbot Loop

Source: https://easydel.readthedocs.io/en/latest/_sources/esurge_examples.rst

Demonstrates the main execution loop for interacting with the ChatBot instance.

```python
if __name__ == "__main__":
        bot = ChatBot()

        print("ChatBot initialized. Type 'quit' to exit, 'reset' to clear history.")

        while True:
            user_input = input("\nYou: ")

            if user_input.lower() == 'quit':
                break
            elif user_input.lower() == 'reset':
                bot.reset()
                print("Conversation reset.")
                continue

            bot.chat(user_input)
```

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

### JSONL format example

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/sources.md

JSONL format, recommended for large datasets, where each line is a self-contained JSON object.

```jsonl
{"text": "First example"}
{"text": "Second example"}
{"text": "Third example"}
```

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

### Configure Model

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/trainer_protocol.md

Implement this abstract method to load your model, create an optimizer and scheduler, and build the initial `EasyDeLState`.

```python
    @abstractmethod
    def configure_model(self) -> TrainerConfigureModelOutput:
        # Load model, create optimizer/scheduler, build EasyDeLState
        # model = ...
        # state = ...
        # self.optimizer = state.opt_state # Or however it's stored
        # self.scheduler = ...
        # return TrainerConfigureModelOutput(model=model, state=state)
        raise NotImplementedError
```

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

### Field Operations

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/transforms.md

Transforms for manipulating fields within data examples, such as selecting or dropping specific keys.

```APIDOC
## SelectFields

### Description
Keeps only the specified fields in the example dictionary.

## DropFields

### Description
Removes the specified fields from the example dictionary.
```

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

### Configure DPOTrainer with DPOConfig

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/dpo.md

Initialize DPOConfig with key parameters for model, training, loss function, sequence length, and optimization.

```python
from easydel.trainers import DPOConfig

dpo_config = DPOConfig(
    # Model and training basics
    model_name="DPOTrainer",  # Name of the model
    learning_rate=1e-6,       # Learning rate for optimization

    # Beta parameter controls how strongly to optimize preferences
    beta=0.1,                 # Temperature parameter for deviation from reference model

    # Loss function options
    loss_type="sigmoid",      # Loss type (sigmoid, hinge, ipo, etc.)
    label_smoothing=0.0,      # Smoothing factor for labels

    # Sequence length parameters
    max_length=512,           # Maximum total sequence length
    max_prompt_length=256,    # Maximum length for prompts
    max_completion_length=256,# Maximum length for completions

    # Reference model control
    reference_free=False,     # Whether to use reference-free variant
    sync_ref_model=False,     # Periodically sync reference model
    ref_model_sync_steps=64,  # Steps between reference model syncs

    # Training optimization
    disable_dropout=True,     # Disable dropout during training
    total_batch_size=16,      # Total batch size
    gradient_accumulation_steps=1  # Steps for gradient accumulation
)
```

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

### Use ShardedDataSource

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/trainer_integration.md

Demonstrates loading data from local or cloud storage using ShardedDataSource for trainers.

```python
from easydel.data import ParquetShardedSource, JsonShardedSource

# From local files
source = ParquetShardedSource("data/*.parquet")

# From cloud
source = ParquetShardedSource(
    "gs://bucket/data/*.parquet",
    storage_options={"token": "cloud"},
)

# Use with any trainer
trainer = ed.SFTTrainer(
    model=model,
    train_dataset=source,  # Works directly
    processing_class=tokenizer,
    arguments=ed.SFTConfig(...),
)
```

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

### Launch eSurge API Server

Source: https://easydel.readthedocs.io/en/latest/_sources/esurge.rst

Starts an OpenAI-compatible API server with optional API key enforcement and monitoring.

```python
# Initialize engine (as shown above)
engine = ed.eSurge(model=model, tokenizer=tokenizer, ...)

# Start monitoring (optional)
engine.start_monitoring()

# Launch API server (API keys optional)
server = ed.eSurgeApiServer(
	engine,
	require_api_key=True,                 # Optional: enforce keys at runtime
	api_keys={"demo-key": {"label": "UI"}},  # Pre-provisioned keys
)

# Additional keys can be issued programmatically
print("Generated key:", server.generate_api_key(label="notebook"))

server.fire(host="0.0.0.0", port=8000)
```

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

### Load Custom Evaluation Dataset

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/sft.md

Example of loading a separate dataset for evaluation during the fine-tuning process using the `load_dataset` function.

```python
# Load evaluation dataset
eval_dataset = load_dataset(
    "Anthropic/hh-rlhf",
    data_dir="helpful-base",
    split="validation[:10%]"
)
```

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

### Command Line SFT Training

Source: https://easydel.readthedocs.io/en/latest/_sources/trainers/sft.md

Execute supervised fine-tuning directly from the command line using the easydel.scripts.finetune.sft module. Specify model, dataset, and training parameters.

```bash
python -m easydel.scripts.finetune.sft \
  --repo_id meta-llama/Llama-3.1-8B-Instruct \
  --dataset_name trl-lib/Capybara \
  --dataset_split "train" \
  --dataset_text_field messages \
  --attn_mechanism vanilla \
  --max_sequence_length 2048 \
  --packing True \
  --total_batch_size 16 \
  --learning_rate 2e-5 \
  --learning_rate_end 5e-6 \
  --num_train_epochs 3 \
  --do_last_save \
  --save_steps 1000 \
  --use_wandb
```

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

### Whisper API Response Formats

Source: https://easydel.readthedocs.io/en/latest/_sources/multimodality/audio_language.md

Examples of the JSON response structures returned by the API for standard text and timestamped segments.

```json
{
  "text": "This is the transcribed text from the audio file."
}
```

```json
{
  "text": "This is the transcribed text with timestamps.",
  "segments": [
    {
      "id": 0,
      "start": 0.0,
      "end": 3.5,
      "text": "This is the first segment"
    },
    {
      "id": 1,
      "start": 3.5,
      "end": 7.2,
      "text": "This is the second segment"
    }
  ]
}
```

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

### Configure eopod Project

Source: https://easydel.readthedocs.io/en/latest/_sources/trc-welcome.md

Sets up the TPU project, zone, and name for the eopod environment.

```shell
eopod configure --project-id YOUR_PROJECT_ID --zone YOUR_ZONE --tpu-name YOUR_TPU_NAME
```

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

### Create and Chain Pipeline Stages

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/pipeline.md

Demonstrates creating a pipeline from configuration and chaining various processing stages. Use this to build a complete data processing workflow.

```python
from easydel.data import Pipeline, PipelineConfig, DatasetConfig

# Create pipeline
config = PipelineConfig(datasets=[DatasetConfig(data_files="data/*.parquet")])
pipeline = Pipeline.from_config(config)

# Chain stages
result = (
    pipeline
    .source()      # Load data sources
    .tokenize()    # Apply tokenization
    .mix()         # Mix multiple datasets
    .pack()        # Pack sequences
    .load()        # Create data loader
    .build()       # Get final iterator
)

# Use in training
for batch in result:
    train_step(batch)
```

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

### Run Whisper Server via CLI

Source: https://easydel.readthedocs.io/en/latest/_sources/whisper_api.md

Start the Whisper server using the built-in CLI. Specify the model and port.

```bash
python -m easydel.inference.vwhisper.cli --model "openai/whisper-large-v3-turbo" --port 8000
```

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

### Initialize KTOTrainer

Source: https://easydel.readthedocs.io/en/latest/_sources/easydata/trainer_integration.md

Configures and trains a KTO model using paired preference data.

```python
dataset = load_dataset("trl-lib/ultrafeedback_binarized", split="train")

trainer = ed.KTOTrainer(
    model=policy_model,
    reference_model=ref_model,
    train_dataset=dataset,
    processing_class=tokenizer,
    arguments=ed.KTOConfig(
        max_prompt_length=512,
        max_completion_length=512,
        beta=0.1,
        loss_type="kto",
    ),
)
trainer.train()
```

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

### Create Chat Engine

Source: https://easydel.readthedocs.io/en/latest/esurge.html

Initializes an eSurge engine for chat applications. Ensure the model ID is valid and necessary dependencies are installed.

```python
import os
os.environ["EASYDEL_AUTO"] = "1"

from jax import lax, numpy as jnp
from transformers import AutoTokenizer
import easydel as ed

def create_chat_engine(model_id="microsoft/phi-2"):
    """Create an eSurge engine for chat."""
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    tokenizer.pad_token_id = tokenizer.eos_token_id

    model = ed.AutoEasyDeLModelForCausalLM.from_pretrained(
        model_id,
        dtype=jnp.bfloat16,
        param_dtype=jnp.bfloat16,
        auto_shard_model=True,
        config_kwargs=ed.EasyDeLBaseConfigDict(
            attn_mechanism=ed.AttentionMechanisms.RAGGED_PAGE_ATTENTION_V2,
        ),
    )

    engine = ed.eSurge(
        model=model,
        tokenizer=tokenizer,
        max_model_len=4096,
        max_num_seqs=16,
        hbm_utilization=0.9,
        page_size=64,
        esurge_name="chat-engine",
    )

    return engine
```

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

### Display Script Help

Source: https://easydel.readthedocs.io/en/latest/_sources/trc-welcome.md

Shows available parameters for a specific fine-tuning script.

```shell
eopod run python -m easydel.scripts.finetune.dpo --help
```

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

### Define Dataset Formats

Source: https://easydel.readthedocs.io/en/latest/trainers/sft.html

Examples of supported data structures for training, including simple text, instruction-response pairs, and conversation lists.

```json
{
  "text": "This is a training example for the language model."
}
```

```json
{
  "instruction": "Write a short poem about machine learning.",
  "response": "Silicon minds learn and grow,\nPatterns found in data's flow.\nMathematical art so precise,\nLearning once, then twice, then thrice."
}
```

```json
{
  "messages": [
    {"role": "system", "content": "You are a helpful AI assistant."},
    {"role": "user", "content": "What is machine learning?"},
    {"role": "assistant", "content": "Machine learning is a subset of artificial intelligence..."}
  ]
}
```

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

### Build a Data Processing Pipeline

Source: https://easydel.readthedocs.io/en/latest/easydata/pipeline.html

Demonstrates creating a pipeline from configuration and chaining various processing stages. Use this to define a sequence of data transformations.

```python
from easydel.data import Pipeline, PipelineConfig, DatasetConfig

# Create pipeline
config = PipelineConfig(datasets=[DatasetConfig(data_files="data/*.parquet")])
pipeline = Pipeline.from_config(config)

# Chain stages
result = (
    pipeline
    .source()      # Load data sources
    .tokenize()    # Apply tokenization
    .mix()         # Mix multiple datasets
    .pack()        # Pack sequences
    .load()        # Create data loader
    .build()       # Get final iterator
)

# Use in training
for batch in result:
    train_step(batch)

```
Tokens:
Raw