### Setup Environment using Conda

Source: https://github.com/alexiglad/ebt/blob/main/README.md

Instructions for setting up the project environment using Conda. This involves creating a new environment, activating it, and installing dependencies from a requirements file. It also mentions optional environment variables like HF_HOME and HF_TOKEN.

```bash
conda create -n ebt python=3.12
conda activate ebt
pip install -r requirements.txt
```

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

### Install FFmpeg via APT (Linux)

Source: https://github.com/alexiglad/ebt/blob/main/data/vid/README.md

This command installs the FFmpeg package, which includes the ffprobe utility, on Debian-based Linux systems using the APT package manager. FFprobe is used for reading video durations.

```bash
apt install ffmpeg
```

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

### Install FFmpeg via Conda

Source: https://github.com/alexiglad/ebt/blob/main/data/vid/README.md

This command installs the FFmpeg package, including ffprobe, within a Conda environment using the conda-forge channel. This is an alternative installation method for systems using Conda.

```bash
conda install -c conda-forge ffmpeg
```

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

### Run Minimalistic NLP Training Loop

Source: https://github.com/alexiglad/ebt/blob/main/README.md

Executes an example minimalistic training loop for Transformer++ vs EBT language models. Note that this example does not reproduce the paper's results but serves as a basic demonstration.

```python
python example_code/minimal_nlp_training_loop.py
```

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

### Run Bash Script Directly

Source: https://github.com/alexiglad/ebt/blob/main/README.md

A quick start method to run a job script directly using bash. This is a straightforward way to execute pretraining or other job configurations.

```bash
bash job_scripts/nlp/pretrain/ebt_s1.sh
```

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

### Initialize Baseline Transformer NLP Model

Source: https://context7.com/alexiglad/ebt/llms.txt

Initializes a Baseline_Transformer_NLP model with specified hyperparameters and performs a forward pass to obtain log probabilities. This snippet demonstrates the basic setup and usage of the transformer model for language modeling tasks.

```python
hparams = {
    'tokenizer': 'EleutherAI/gpt-neox-20b',
    'embedding_dim': 384,
    'num_transformer_blocks': 6,
    'multiheaded_attention_heads': 6,
    'weight_initialization_method': 'xavier',
    'weight_initialization_gain': 1.0,
    'execution_mode': 'pretrain'
}

model = Baseline_Transformer_NLP(hparams)

# Forward pass returns log probabilities
input_ids = torch.randint(0, 50277, (2, 256))  # batch_size=2, seq_len=256
log_probs = model(input_ids, learning=True)

# log_probs: [B*S, V] log probabilities over vocabulary
print(f"Output shape: {log_probs.shape}")  # torch.Size([512, 50277])
```

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

### Run Bash Script using SLURM Executor

Source: https://github.com/alexiglad/ebt/blob/main/README.md

Recommended method for running bash scripts on High-Performance Computing (HPC) clusters with SLURM installed. This approach uses a `slurm_executor.sh` script for modularity and includes a mandatory parameter for SLURM script configuration.

```bash
bash slurm_executor.sh reference_a100 job_scripts/nlp/pretrain/ebt_s1.sh
```

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

### Loading Pretrained EBT Model Checkpoints

Source: https://context7.com/alexiglad/ebt/llms.txt

Python code demonstrating two methods for loading pretrained EBT model checkpoints. Method 1 shows loading through a training script by setting specific hyperparameters like `only_test_model_ckpt` and `only_test`. Method 2 demonstrates direct loading using PyTorch Lightning's `load_from_checkpoint` method.

```python
from model.model_utils import load_trained_pl_model
import pytorch_lightning as pl

# Load checkpoint with automatic hparam restoration
checkpoint_path = "checkpoints/ebt-xxs-epoch=10-step=211000.ckpt"

# Method 1: Load through training script
hparams = {
    'only_test_model_ckpt': checkpoint_path,
    'only_test': True,
    'execution_mode': 'inference',
    # Override specific inference params
    'infer_max_gen_len': 128,
    'infer_temp': 0.8,
    'infer_topp': 0.9,
}
# Model hparams automatically inherited from checkpoint

# Method 2: Direct PyTorch Lightning load
from base_model_trainer import ModelTrainer
model = ModelTrainer.load_from_checkpoint(checkpoint_path)
model.eval()

```

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

### Execute Slurm Jobs for Training

Source: https://context7.com/alexiglad/ebt/llms.txt

Demonstrates how to execute training jobs on HPC clusters using Slurm. It covers direct execution of job scripts, using the `slurm_executor.sh` script, and configuring Slurm header parameters. Multi-node training with `srun` is also illustrated.

```bash
# Execute job script directly
bash job_scripts/nlp/pretrain/ebt_s1.sh

# Execute via slurm executor (recommended for HPC)
bash slurm_executor.sh reference_a100 job_scripts/nlp/pretrain/ebt_s1.sh

# Slurm script header configuration (job_scripts/slurm_headers/reference_a100.slurm)
# #!/bin/bash
# #SBATCH --partition=gpu
# #SBATCH --nodes=1
# #SBATCH --gpus-per-node=4
# #SBATCH --ntasks-per-node=4
# #SBATCH --mem=256G
# #SBATCH --time=48:00:00

# Multi-node training with srun
bash job_scripts/nlp/pretrain/ebt_s1_mn.sh
# Uses: srun python train_model.py --gpus "-1" --num_nodes 2
```

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

### Unzip SSv2 Labels Archive

Source: https://github.com/alexiglad/ebt/blob/main/data/vid/README.md

This bash command is used to unzip the label files for the Something-something-v2 dataset. It assumes the zip file is present in the current directory.

```bash
unzip 20bn-something-something-download-package-labels.zip
```

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

### Learning Rate Scheduling with Warmup and Cosine Annealing

Source: https://context7.com/alexiglad/ebt/llms.txt

Python code demonstrating how to configure a learning rate scheduler with warmup and cosine annealing using PyTorch. It separates parameters for weight decay, initializes an AdamW optimizer, and sets up a `WarmUpCosineAnnealingLR` scheduler. The code includes a basic training loop with gradient clipping and scheduler steps.

```python
from optimization import WarmUpCosineAnnealingLR, exclude_bias_and_norm
import torch.optim as optim
import torch

# Separate parameters for weight decay
param_groups = exclude_bias_and_norm(
    model,
    weight_decay=0.01,
    skip_list=['alpha', 'langevin_dynamics_noise_std']  # Don't decay special params
)

optimizer = optim.AdamW(param_groups, lr=0.0012, betas=(0.9, 0.95))

# Warmup + cosine annealing scheduler
scheduler = WarmUpCosineAnnealingLR(
    optimizer,
    warm_up_steps=10000,
    max_steps=1000000,
    min_lr_scale=10,  # min_lr = peak_lr / 10
    max_scheduling_steps=1000000
)

# Training loop
for step in range(1000000):
    loss = model.training_step(batch, step)
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    optimizer.step()
    scheduler.step()

    if step % 100 == 0:
        print(f"Step {step}, LR: {scheduler.get_last_lr()[0]:.6f}, Loss: {loss.item():.4f}")
```

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

### Concatenate and Unzip SSv2 Archive

Source: https://github.com/alexiglad/ebt/blob/main/data/vid/README.md

This code snippet demonstrates how to concatenate multiple parts of the Something-something-v2 dataset archive into a single file and then extract its contents. It's a bash command sequence for processing downloaded archive files.

```bash
cat 20bn-something-something-v2-* > ssv2_archive
tar -xvf ssv2_archive
```

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

### Distributed Training Configuration with PyTorch Lightning

Source: https://context7.com/alexiglad/ebt/llms.txt

Python code for configuring distributed training using PyTorch Lightning and the DDP strategy. It initializes a `ModelTrainer` with hyperparameters and then sets up a `pl.Trainer` instance, specifying devices, number of nodes, DDP strategy parameters, gradient accumulation, mixed precision, and logging intervals. The trainer is then used to fit the model.

```python
import pytorch_lightning as pl
from pytorch_lightning.strategies import DDPStrategy
from base_model_trainer import ModelTrainer

# Initialize model trainer with hyperparameters
hparams = {
    'model_name': 'ebt',
    'model_size': 's',
    'modality': 'NLP',
    'batch_size_per_device': 32,
    'accumulate_grad_batches': 2,
    'peak_learning_rate': 0.0012,
    'max_steps': 1000000,
    'gpus': -1,  # Use all available GPUs
    'num_nodes': 1,
    'context_length': 256,
    'dataset_name': 'pajama',
    # ... other hparams
}

model_trainer = ModelTrainer(hparams)

# Configure trainer with DDP
trainer = pl.Trainer(
    max_steps=hparams['max_steps'],
    devices=-1,  # All GPUs
    num_nodes=hparams['num_nodes'],
    strategy=DDPStrategy(
        find_unused_parameters=False,
        gradient_as_bucket_view=True
    ),
    accumulate_grad_batches=hparams['accumulate_grad_batches'],
    gradient_clip_val=1.0,
    precision='16-mixed',  # Mixed precision training
    log_every_n_steps=50,
    val_check_interval=15000,
)

trainer.fit(model_trainer)
# Automatically handles gradient synchronization across GPUs
```

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

### Minimal PyTorch Lightning Training Loop for EBT

Source: https://context7.com/alexiglad/ebt/llms.txt

Implements a minimal training loop using PyTorch Lightning to compare EBT against a baseline transformer. It defines a `ModelWrapper` class handling dataset loading, training steps, and optimizer configuration, suitable for pre-training tasks.

```python
import torch
import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
from torch.utils.data import DataLoader
from data.nlp.pajama_dataloader import RedPajamaDataset
from data.nlp.collator import NLP_HF_Collator
from model.nlp.baseline_transformer import Baseline_Transformer_NLP
from model.nlp.ebt import EBT_NLP

class ModelWrapper(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.save_hyperparameters(hparams)

        model_cls = {
            "baseline_transformer": Baseline_Transformer_NLP,
            "ebt": EBT_NLP,
        }[self.hparams.model_name]

        self.model = model_cls(self.hparams)
        self.dataset = RedPajamaDataset(self.hparams)
        self.collate_fn = NLP_HF_Collator(self.hparams)

    def training_step(self, batch, batch_idx):
        metrics = self.model.forward_loss_wrapper(batch, "train")
        loss = metrics["loss"]
        self.log_dict({f"train_{k}": v for k, v in metrics.items()}, prog_bar=True)
        return loss

    def configure_optimizers(self):
        return torch.optim.AdamW(self.model.parameters(), lr=self.hparams.lr)

    def train_dataloader(self):
        workers = torch.cuda.device_count() * self.hparams.num_workers_per_gpu
        return DataLoader(
            self.dataset,
            batch_size=self.hparams.batch_size_per_device,
            shuffle=True,
            num_workers=workers,
            collate_fn=self.collate_fn,
        )

# Train EBT model
hparams = {
    'lr': 1e-3,
    'batch_size_per_device': 32,
    'num_workers_per_gpu': 12,
    'max_steps': 100000,
    'dataset_dir': '',
    'dataset_name': 'pajama',
    'context_length': 256,
    'pretokenize_dataset': True,
    'tokenizer': 'EleutherAI/gpt-neox-20b',
    'model_name': 'ebt',
    'embedding_dim': 384,
    'num_transformer_blocks': 6,
    'multiheaded_attention_heads': 6,
    'ffn_dim_multiplier': 1,
    'weight_initialization_method': 'xavier',
    'weight_initialization_gain': 1.0,
    'execution_mode': 'pretrain',
    'mcmc_step_size': 500.0,
    'mcmc_step_size_lr_multiplier': 1500.0,
    'mcmc_num_steps': 2,
    'ebt_type': 'time_embed',
    'normalize_initial_condition': True,
    'denoising_initial_condition': 'random_noise',
    'mcmc_step_size_learnable': True,
    'no_mcmc_detach': False,
    'debug_unused_parameters': False
}

model = ModelWrapper(hparams)
logger = WandbLogger(name="ebt_training", project="nlp_pretrain", entity="")

trainer = pl.Trainer(
    max_steps=hparams['max_steps'],
    devices=-1,
    logger=logger,
    enable_model_summary=True,
    enable_checkpointing=True,
)
trainer.fit(model)
```

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

### System 2 EBT with Causal Replay Buffer

Source: https://context7.com/alexiglad/ebt/llms.txt

Python code for initializing an EBT model for NLP tasks with a causal replay buffer, enabling advanced System 2 thinking. Hyperparameters include tokenizer, embedding dimensions, transformer blocks, attention heads, MCMC parameters, replay buffer configuration, and batch size. The code snippet illustrates the initialization and a conceptual training loop where the replay buffer is utilized.

```python
from model.replay_buffer import CausalReplayBuffer
from model.nlp.ebt import EBT_NLP

# Initialize EBT with replay buffer for System 2 thinking
hparams = {
    'tokenizer': 'EleutherAI/gpt-neox-20b',
    'embedding_dim': 768,
    'num_transformer_blocks': 12,
    'multiheaded_attention_heads': 12,
    'mcmc_step_size': 500.0,
    'mcmc_num_steps': 5,  # More steps for System 2
    'ebt_type': 'time_embed',
    'mcmc_step_size_learnable': True,
    'mcmc_replay_buffer': True,
    'mcmc_replay_buffer_size': 10000,
    'mcmc_replay_buffer_sample_bs_percent': 0.5,  # 50% from buffer
    'batch_size_per_device': 32,
    'execution_mode': 'pretrain',
    # ... other S2 params
}

model = EBT_NLP(hparams)

# Replay buffer stores past optimization trajectories
# During training, 50% of batch uses fresh data, 50% from buffer
for batch_idx, batch in enumerate(train_loader):
    metrics = model.forward_loss_wrapper(batch, phase="train")
    # Replay buffer automatically samples and provides logits to forward()
    # This enables the model to learn from past reasoning traces

print(f"Replay buffer size: {model.replay_buffer.current_size}")
```

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

### Train EBT for Video Prediction on Kinetics-400

Source: https://context7.com/alexiglad/ebt/llms.txt

Command-line arguments for training the EBT model for video prediction on the Kinetics-400 dataset. Specifies model architecture, dataset, image dimensions, number of frames, backbone type, MCMC parameters, batch size, learning rate, training steps, GPU usage, worker count, and Weights & Biases project.

```shell
python train_model.py \
  --run_name ebt-vid-k400 \
  --modality "VID" \
  --model_name "ebt_vid" \
  --model_size "s" \
  --dataset_name "k400" \
  --image_dims 128 128 \
  --num_frames 16 \
  --backbone_type "vae" \
  --mcmc_step_size 1.0 \
  --mcmc_num_steps 3 \
  --batch_size_per_device 16 \
  --peak_learning_rate 5e-5 \
  --max_steps 1000000 \
  --gpus "-1" \
  --num_workers 8 \
  --wandb_project "vid_prediction"
```

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

### Load and Pretokenize RedPajama Dataset

Source: https://context7.com/alexiglad/ebt/llms.txt

Loads and pretokenizes the RedPajama-100B dataset for efficient training using the `RedPajamaDataset` class. It requires specifying hyperparameters such as context length, tokenizer, and dataset name. The preprocessed dataset is automatically cached to disk.

```python
from data.nlp.pajama_dataloader import RedPajamaDataset
from transformers import AutoTokenizer

# Initialize dataset with pretokenization
hparams = {
    'execution_mode': 'pretrain',
    'context_length': 256,
    'dataset_dir': '',  # Uses HF_HOME env variable
    'tokenizer': 'EleutherAI/gpt-neox-20b',
    'pretokenize_dataset': True,
    'dataset_name': 'pajama',
    'num_workers': 12,
    'num_gpus': 4,
    'batch_size_per_device': 32
}

dataset = RedPajamaDataset(hparams)

# Access tokenized samples
sample = dataset[0]
print(f"Input IDs shape: {sample['input_ids'].shape}")
print(f"Attention mask shape: {sample['attention_mask'].shape}")
print(f"Dataset length: {len(dataset)}")

# Dataset automatically caches preprocessed version to disk
# Location: $HF_HOME/pajama_preprocessed/EleutherAI_gpt-neox-20b/max_length_257/
```

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

### Run Inference on Language Understanding Benchmarks

Source: https://context7.com/alexiglad/ebt/llms.txt

Executes inference for language understanding tasks using the EBT model. This script requires specifying model details, tokenizer, dataset, and inference parameters. It supports various configurations for context length, batch size, and generation strategies.

```bash
python train_model.py \
  --run_name ebt-xxs-inference-lambada \
  --modality "NLP" \
  --model_name "ebt" \
  --model_size "xxs" \
  --tokenizer "EleutherAI/gpt-neox-20b" \
  --normalize_initial_condition \
  --ebt_type "time_embed" \
  --denoising_initial_condition "random_noise" \
  --context_length 256 \
  --gpus "-1" \
  --batch_size_per_device 8 \
  --dataset_name "lambada" \
  --num_workers 12 \
  --wandb_project "nlp_inference_accuracy" \
  --execution_mode "inference" \
  --infer_ebt_advanced \
  --infer_langevin_dynamics_noise 1.0 \
  --infer_ebt_num_steps 2 \
  --only_test \
  --only_test_model_ckpt "checkpoints/ebt-xxs-epoch=10.ckpt" \
  --infer_max_gen_len 128 \
  --infer_topp 0.9 \
  --infer_temp 0.8 \
  --set_matmul_precision "medium"
```

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

### Configure Model Size Presets

Source: https://context7.com/alexiglad/ebt/llms.txt

Automatically configures model dimensions based on predefined size presets (xxs, xs, s, m, l). This script demonstrates how to parse command-line arguments for model size and dynamically set the number of transformer blocks, attention heads, and embedding dimension.

```python
from model.model_utils import model_sizes

# Available model sizes with automatic parameter configuration
# model_sizes = {
#     "xxs": {'num_transformer_blocks': 6, 'multiheaded_attention_heads': 6, 'embedding_dim': 384},
#     "xs": {'num_transformer_blocks': 12, 'multiheaded_attention_heads': 12, 'embedding_dim': 768},
#     "s": {'num_transformer_blocks': 24, 'multiheaded_attention_heads': 16, 'embedding_dim': 1024},
#     "m": {'num_transformer_blocks': 32, 'multiheaded_attention_heads': 32, 'embedding_dim': 2048},
#     "l": {'num_transformer_blocks': 48, 'multiheaded_attention_heads': 48, 'embedding_dim': 4096}
# }

# Use in training script
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--model_size', type=str, default='xxs')
args = parser.parse_args(['--model_size', 's'])

if args.model_size:
    args.num_transformer_blocks = model_sizes[args.model_size]['num_transformer_blocks']
    args.multiheaded_attention_heads = model_sizes[args.model_size]['multiheaded_attention_heads']
    args.embedding_dim = model_sizes[args.model_size]['embedding_dim']
    print(f"Model: {args.num_transformer_blocks} layers, {args.embedding_dim} dim")
    # Output: Model: 24 layers, 1024 dim
```

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

### EBT Model Forward Pass with Energy Optimization (Python)

Source: https://context7.com/alexiglad/ebt/llms.txt

Demonstrates a forward pass of the EBT_NLP model with energy-based optimization. It initializes the model with specified hyperparameters and performs iterative gradient descent steps to refine predictions, printing the final distribution shape and energy progression.

```python
from model.nlp.ebt import EBT_NLP
import torch

# Initialize EBT model
hparams = {
    'tokenizer': 'EleutherAI/gpt-neox-20b',
    'embedding_dim': 384,
    'num_transformer_blocks': 6,
    'multiheaded_attention_heads': 6,
    'mcmc_step_size': 500.0,
    'mcmc_step_size_lr_multiplier': 1500.0,
    'mcmc_num_steps': 2,
    'ebt_type': 'time_embed',
    'normalize_initial_condition': True,
    'denoising_initial_condition': 'random_noise',
    'mcmc_step_size_learnable': True,
    'no_mcmc_detach': False,
    'weight_initialization_method': 'xavier',
    'weight_initialization_gain': 1.0,
    'execution_mode': 'pretrain',
    'debug_unused_parameters': False
}

model = EBT_NLP(hparams)

# Forward pass with energy optimization
input_ids = torch.randint(0, 50277, (2, 256))  # batch_size=2, seq_len=256
predicted_distributions, predicted_energies = model(input_ids, learning=True)

# predicted_distributions: list of [B*S, V] tensors for each MCMC step
# predicted_energies: list of energy values showing optimization progress
print(f"Final distribution shape: {predicted_distributions[-1].shape}")
print(f"Energy progression: {[e.mean().item() for e in predicted_energies]}")
```

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

### Train EBT Language Model (Bash)

Source: https://context7.com/alexiglad/ebt/llms.txt

Trains an Energy-Based Transformer (EBT) for language modeling using pretokenized datasets. This script configures System 1 EBT with 2 MCMC steps on the RedPajama dataset, specifying various model and training hyperparameters.

```bash
python train_model.py \
  --run_name ebt-xxs-bs_256_s1 \
  --modality "NLP" \
  --model_name "ebt" \
  --model_size "xxs" \
  --pretokenize_dataset \
  --tokenizer "EleutherAI/gpt-neox-20b" \
  --normalize_initial_condition \
  --ebt_type "time_embed" \
  --denoising_initial_condition "random_noise" \
  --mcmc_step_size_learnable \
  --mcmc_step_size 500.0 \
  --mcmc_step_size_lr_multiplier 1500.0 \
  --mcmc_num_steps 2 \
  --context_length 256 \
  --gpus "-1" \
  --peak_learning_rate 0.0012 \
  --batch_size_per_device 32 \
  --accumulate_grad_batches 2 \
  --gradient_clip_val 1.0 \
  --weight_decay 0.01 \
  --min_lr_scale 10 \
  --max_steps 1000000 \
  --warm_up_steps 10000 \
  --dataset_name "pajama" \
  --num_workers 12 \
  --validation_split_pct 0.0005 \
  --val_check_interval 15000 \
  --wandb_project "nlp_pretrain" \
  --log_model_archi \
  --log_gradients \
  --set_matmul_precision "medium"
```

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

### Train Baseline Transformer Model (Bash)

Source: https://context7.com/alexiglad/ebt/llms.txt

Trains a standard feed-forward transformer model for comparison purposes. This bash script configures a baseline transformer for language modeling on the RedPajama dataset, setting relevant training parameters.

```bash
python train_model.py \
  --run_name baseline-xxs-bs_256 \
  --modality "NLP" \
  --model_name "baseline_transformer" \
  --model_size "xxs" \
  --pretokenize_dataset \
  --tokenizer "EleutherAI/gpt-neox-20b" \
  --context_length 256 \
  --gpus "-1" \
  --peak_learning_rate 0.0012 \
  --batch_size_per_device 32 \
  --accumulate_grad_batches 2 \
  --gradient_clip_val 1.0 \
  --weight_decay 0.01 \
  --max_steps 1000000 \
  --warm_up_steps 10000 \
  --dataset_name "pajama" \
  --num_workers 12 \
  --wandb_project "nlp_pretrain"
```

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

### Train EBT for Image Denoising and Generation

Source: https://context7.com/alexiglad/ebt/llms.txt

Trains an Energy-Based Transform (EBT) model for image denoising and generation. This script configures the model for image modality, specifies image dimensions, backbone type, and MCMC parameters. It also includes parameters for training duration and device allocation.

```bash
# Train EBT for image denoising (System 1)
python train_model.py \
  --run_name ebt-img-s1-denoise \
  --modality "IMG" \
  --model_name "ebt_denoise" \
  --model_size "s" \
  --dataset_name "ImageNet1k" \
  --image_dims 64 64 \
  --backbone_type "vae" \
  --vae_normalization \
  --mcmc_step_size 0.1 \
  --mcmc_num_steps 2 \
  --ebt_type "adaln" \
  --batch_size_per_device 64 \
  --peak_learning_rate 1e-4 \
  --max_steps 500000 \
  --gpus "-1" \
  --wandb_project "img_generation"
```

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

### Generate Images from Trained Checkpoint

Source: https://context7.com/alexiglad/ebt/llms.txt

Generates images using a trained EBT denoising model. This script sets the execution mode to 'inference', specifies the checkpoint path, and configures the number of samples to generate and the GPU to use. It's used for evaluating the generative capabilities of the trained model.

```bash
# Generate images from trained checkpoint
python train_model.py \
  --run_name ebt-img-inference \
  --modality "IMG" \
  --model_name "ebt_denoise" \
  --model_size "s" \
  --execution_mode "inference" \
  --only_test \
  --only_test_model_ckpt "checkpoints/ebt-img-s1-epoch=50.ckpt" \
  --infer_num_samples 100 \
  --gpus "1"
```

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

### Load Model Weights in PyTorch

Source: https://context7.com/alexiglad/ebt/llms.txt

Loads model weights from a checkpoint file and populates the model's state dictionary. This is useful for resuming training or performing inference with a pre-trained model. Requires a PyTorch model and a checkpoint dictionary containing 'state_dict' and training-related information like 'global_step' and 'callbacks'.

```python
checkpoint = torch.load(checkpoint_path)
model.load_state_dict(checkpoint['state_dict'])

print(f"Loaded model from step {checkpoint['global_step']}")
print(f"Training loss: {checkpoint['callbacks']['loss']:.4f}")
```

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

### Baseline Transformer Forward Pass (Python)

Source: https://context7.com/alexiglad/ebt/llms.txt

Illustrates a standard autoregressive transformer prediction without energy optimization. This Python snippet shows the import statement for the Baseline_Transformer_NLP model, setting up for a forward pass.

```python
from model.nlp.baseline_transformer import Baseline_Transformer_NLP
import torch


```

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

### Compute Language Modeling Loss and Perplexity

Source: https://context7.com/alexiglad/ebt/llms.txt

Calculates the cross-entropy loss and perplexity for next-token prediction using the EBT model. This involves preparing input batches, obtaining predicted distributions, and applying PyTorch's Negative Log Likelihood loss function.

```python
import torch
import torch.nn.functional as F

# Prepare batch with input_ids of shape [batch_size, seq_len+1]
batch = {
    'input_ids': torch.randint(0, 50277, (8, 257)).unsqueeze(1)  # [8, 1, 257]
}

# EBT loss calculation
input_ids = batch['input_ids'].squeeze(dim=1)[:, :-1]  # [8, 256]
predicted_distributions = model.forward_loss_wrapper(batch, phase="train")

# Compute NLL loss
next_tokens = batch['input_ids'].squeeze(dim=1)[:, 1:]  # [8, 256]
next_tokens = next_tokens.reshape(-1)  # [2048]
cce_loss = F.nll_loss(predicted_distributions[-1], next_tokens, ignore_index=0)
perplexity = torch.exp(cce_loss)

print(f"Loss: {cce_loss.item():.4f}, Perplexity: {perplexity.item():.4f}")
```

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