Training with PyTorch Lightning: Structured MLOps Development

In this lesson, you will learn how to build a fully modular, reproducible, and production-friendly training pipeline using PyTorch Lightning and Hydra. We will train a sentiment-classification model with a clean MLOps-ready structure that scales as your projects grow.

This lesson is the 1st in a 2-part series on PyTorch Lightning:

1. Training with PyTorch Lightning: Structured MLOps Development (this tutorial)
2. Lesson 2

To learn how to structure deep learning code for reliable, maintainable training workflows, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Why PyTorch Lightning Improves Reproducible MLOps Training Pipelines

Deep learning projects begin as small experiments (e.g., a single script, a few functions, and a quick training loop). But as your model grows, your dataset expands, and your experimentation increases, that simple script becomes a bottleneck. It collects data loading, model code, training logic, metrics, logging, and CLI arguments all in one place. This makes the code harder to debug, harder to reproduce, and nearly impossible to scale.

From an MLOps perspective, this is a critical issue. Reproducibility, modularity, logging, version control, and collaboration all depend on a clean separation of concerns. When data, modeling, and training logic live inside one file, versioning becomes fragile. Dockerization becomes harder. Running distributed or mixed-precision training becomes messy. And experiment tracking ends up inconsistent.

This is exactly where PyTorch Lightning shines. Instead of mixing everything together, Lightning organizes your training workflow into well-defined components:

• LightningModule: contains your model and training/validation/test logic
• LightningDataModule: manages data downloading, tokenization, and data loaders
• Trainer: handles the engineering: GPUs, mixed precision, checkpointing, logging, and training loops

You keep your PyTorch code. Lightning handles the boilerplate.

This separation is not just “clean code.” It is an MLOps superpower. With structure in place, you can:

• Version your data, model code, and configs separately
• Run consistent experiments across environments
• Scale from CPU → single GPU → multi-GPU → multi-node without changing code
• Produce reproducible artifacts for DVC and deployment
• Integrate Hydra configs for repeatable pipelines

In this lesson, you will see how Lightning transforms a typical sentiment-classification workflow into a modular, reproducible, production-ready training pipeline (i.e., the kind expected in real-world MLOps and LLMOps teams).

To learn how to structure reproducible ML training with PyTorch Lightning, just keep reading.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Configuring Your Development Environment

Before we begin implementing the training pipeline, let’s configure the development environment. This project uses modern versions of PyTorch, PyTorch Lightning, and the Hugging Face ecosystem, along with Hydra for configuration management. Installing the correct dependencies ensures that your training, data loading, and model export workflows run smoothly.

All the required libraries are pip-installable. Below is the full list of dependencies used in this lesson:

# PyTorch and Lightning
torch>=2.1.0
torchvision>=0.16.0
torchaudio>=2.1.0
pytorch-lightning>=2.1.0

# Hugging Face ecosystem
transformers>=4.35.0
datasets>=2.14.0
tokenizers>=0.15.0

# Configuration management
hydra-core>=1.3.0
omegaconf>=2.3.0

# Metrics and monitoring
torchmetrics>=1.2.0
tensorboard>=2.15.0

# Data processing
numpy>=1.24.0
pandas>=2.0.0

# Utilities
tqdm>=4.66.0
pyyaml>=6.0.0

You can install everything with a single command:

!pip install -r requirements.txt

Your project folder includes a requirements.txt file containing all of the dependencies above. If you are working inside a clean virtual environment (e.g., venv or Conda), this setup will give you everything needed to train the sentiment-classification model in Lesson 1.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Project Structure

Before we write any code, it is important to understand how a real MLOps-ready training workflow is organized. A clean project structure makes your training pipeline easier to maintain, debug, scale across machines, and integrate later with DVC, MLflow, and CI/CD. PyTorch Lightning encourages this modular layout (e.g., through structured components), and Hydra takes care of configuration management in a clean, reproducible way.

Below is the project structure we will use for Lesson 1:

lesson1/
├── configs/
│   ├── config.yaml
│   ├── model/
│   │   └── distilbert.yaml
│   ├── data/
│   │   └── imdb.yaml
│   └── trainer/
│       └── default.yaml
└── src/
    ├── data_module.py
    ├── model_module.py
    ├── train.py
    └── inference.py

The src/ folder contains all the Python logic:

• data_module.py: handles dataset loading, tokenization, and data loaders
• model_module.py: defines the DistilBERT-based classifier
• train.py: orchestrates the training loop using PyTorch Lightning
• inference.py: allows you to test the trained model interactively

The configs/ directory is where we store all configuration files, grouped by purpose (model, data, trainer). This keeps hyperparameters and training settings cleanly separated from code, which is a fundamental principle in MLOps.

Hydra in 60 Seconds (All You Need for This Lesson)

Hydra is a configuration framework from Facebook Research that lets us manage all training settings (e.g., model hyperparameters, data parameters, and trainer options) in clean, reusable YAML files. Instead of hard-coding values inside Python, Hydra loads these configs automatically and lets you override anything from the command line.

This makes your training pipeline more reproducible, easier to maintain, and far more MLOps-friendly.

For example, you can change batch size, learning rate, or precision without editing the code:

python src/train.py data.batch_size=16 model.lr=3e-5 trainer.precision=16-mixed

Hydra merges the YAML configs at runtime (using config.yaml as the root), giving us a clean and scalable way to manage experiments.

Now that we understand how the project is organized and how Hydra helps us control training behavior, let us implement the LightningDataModule and build the first part of our training pipeline.

PyTorch LightningDataModule Explained: Building Efficient Data Pipelines

The LightningDataModule is one of the most important building blocks in PyTorch Lightning. It packages every data-related responsibility (e.g., downloading, preprocessing, tokenizing, and creating data loaders) into a single reusable module. This keeps your training loop clean and preserves a sharp separation between data code and model code, which is essential in any production-ready MLOps workflow.

In sentiment classification, the DataModule encapsulates 3 main tasks:

• loading the IMDB dataset: from Hugging Face
• tokenizing text: using a pretrained Transformer tokenizer
• producing data loaders: for train, validation, and test splits

Let us walk through the implementation step-by-step.

Before we start examining the class, let us begin with the imports that power this component of the pipeline:

import pytorch_lightning as pl
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoTokenizer

These 4 imports define everything the DataModule needs. pytorch_lightning gives us the LightningDataModule base class, which enforces a structured, reproducible way to manage data. load_dataset from Hugging Face handles downloading and preparing the IMDB dataset.

PyTorch’s DataLoader will assemble batches and handle multiprocessing, and AutoTokenizer loads the correct tokenizer for whichever transformer model we choose in the configuration.

Now that the foundations are clear, let us walk through the class itself.

Initializing the DataModule (__init__)

The constructor sets up every configurable part of the data pipeline.

Here is the code block we are examining:

class SentimentDataModule(pl.LightningDataModule):
    def __init__(
        self,
        model_name: str = "distilbert-base-uncased",
        dataset_name: str = "imdb",
        max_length: int = 128,
        batch_size: int = 8,
        num_workers: int = 4,
        **kwargs
    ):
        super().__init__()
        self.save_hyperparameters()
        
        self.model_name = model_name
        self.dataset_name = dataset_name
        self.max_length = max_length
        self.batch_size = batch_size
        self.num_workers = num_workers
        
        self.tokenizer = None
        self.train_dataset = None
        self.val_dataset = None
        self.test_dataset = None

The class inherits from pl.LightningDataModule, which means Lightning expects this object to implement the standard methods (prepare_data, setup, and the data loaders). The constructor receives all runtime-configurable parameters (e.g., which model tokenizer to use, which dataset to load, how long sequences should be, and what batch size to apply). These values are stored as instance attributes, but they are also captured automatically through self.save_hyperparameters(), enabling 2 important MLOps behaviors: experiment reproducibility and checkpoint metadata tracking.

The module initializes empty placeholders for the tokenizer and each dataset split. Nothing is loaded at this point. This keeps initialization fast, makes the DataModule safe for multiprocessing, and follows Lightning’s recommended pattern where all heavyweight work happens in the next 2 methods: prepare_data() and setup().

prepare_data(): One-Time Download Step

Next is the method prepare_data() that Lightning calls exactly once on a single process:

The constructor receives all hyperparameters needed to configure data preprocessing. Hydra will pass these values at runtime.

def prepare_data(self):
    load_dataset(self.dataset_name)
    AutoTokenizer.from_pretrained(self.model_name)

prepare_data() is responsible only for downloading the dataset and the tokenizer files. It does not tokenize or process anything; that comes later. The logic here intentionally has no side effects and does not assign anything to self. Lightning enforces this behavior because, in distributed training, only one GPU should perform downloads, while every GPU should independently run tokenization and data setup.

This separation is crucial for large-scale MLOps workflows because it prevents race conditions, avoids repetitive downloads, and maintains determinism across runs regardless of device or environment.

The setup() Method: Tokenization and Dataset Preparation

After the one-time download step in prepare_data(), Lightning calls the setup() method on every GPU (or CPU worker) participating in training. This is where the real work happens: tokenization, dataset formatting, and creating the train/validation/test splits. Because each device executes this method independently, the logic inside must be deterministic and side-effect-free, which is a core requirement when building reproducible training pipelines in distributed environments.

Here is the method we are explaining:

def setup(self, stage=None):
    self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
    
    if stage == "fit" or stage is None:
        dataset = load_dataset(self.dataset_name)

        self.train_dataset = dataset["train"].map(
            self._tokenize_function,
            batched=True,
            remove_columns=dataset["train"].column_names,
            desc="Tokenizing train set"
        )
        self.train_dataset.set_format("torch")

        self.val_dataset = dataset["test"].map(
            self._tokenize_function,
            batched=True,
            remove_columns=dataset["test"].column_names,
            desc="Tokenizing validation set"
        )
        self.val_dataset.set_format("torch")
    
    if stage == "test" or stage is None:
        dataset = load_dataset(self.dataset_name)
        self.test_dataset = dataset["test"].map(
            self._tokenize_function,
            batched=True,
            remove_columns=dataset["test"].column_names,
            desc="Tokenizing test set"
        )
        self.test_dataset.set_format("torch")

The first line initializes a fresh tokenizer instance. Lightning may spawn multiple processes in Distributed Data Parallel (DDP) or Fully Sharded Data Parallel (FSDP) mode, so each process needs its own tokenizer object. Tokenizers are stateless, lightweight, and safe to re-create, making them ideal for device-local initialization.

Lightning passes a stage identifier depending on what it is doing:

• fit: when calling trainer.fit()
• validate: when calling trainer.validate()
• test: when calling trainer.test()
• predict: when calling trainer.predict()
• None: Lightning did not specify a stage (certain internal flows)

Handling the None stage ensures the DataModule works correctly in all cases.

Inside the fit stage, we load the Hugging Face dataset. IMDB contains only train and test splits. Since it does not ship with an official validation set, this lesson uses the test split as validation. (In a real production MLOps pipeline, you would do a train/validation split yourself. This note will be added in the blog.)

Next comes tokenization. Hugging Face Datasets operates like a high-performance dataframe.

.map() applies _tokenize_function to every example or batch, efficiently processing the dataset:

• processes data in batches automatically
• parallelizes when possible
• avoids loading everything into memory

This is far more efficient and cleaner than writing custom loops.

After tokenization, we call self.train_dataset.set_format("torch"). This converts the dataset to output:

• torch.Tensor for inputs
• torch.Tensor for labels

which is exactly what PyTorch Lightning and the data loader expect. A nearly identical process prepares the validation dataset.

The test stage works the same way. Lightning may call fit and test separately, so we load and tokenize test data inside its own conditional block.

The logic inside setup() must be safe to run:

• once per GPU
• once per CPU worker
• once per stage

This is why nothing is written to disk or downloaded in this method. Everything is device-local: load → tokenize → return.

This is essential for:

• reproducibility
• distributed training
• multi-node training
• deterministic behavior in automated pipelines

Tokenization Logic: The Heart of Text Preprocessing

Tokenization is the most important preprocessing step in any natural language processing (NLP) pipeline. It converts raw text into numerical tensors that a Transformer model like DistilBERT can understand. Unlike image data (where pixel arrays are already numeric), text needs to be transformed from variable-length strings into fixed-length integer sequences.

In our DataModule, all tokenization runs through a single helper method:

def _tokenize_function(self, examples):
    tokenized = self.tokenizer(
        examples["text"],
        truncation=True,
        padding="max_length",
        max_length=self.max_length,
    )
    tokenized["labels"] = examples["label"]
    return tokenized

Let us break this down step-by-step.

The _tokenize_function handles the core preprocessing step for every text sample in the dataset. Hugging Face Datasets calls this function in batched mode, so examples["text"] is a list of movie reviews rather than a single string. The tokenizer processes the entire batch at once, which is significantly faster than looping over rows individually and also plays nicely with multiprocessing when multiple workers are specified in the data loader.

Inside the tokenizer call, we enable truncation and padding to a fixed max_length. Truncation ensures that very long reviews do not exceed the model’s maximum input window, keeping GPU memory usage predictable. Padding ensures that shorter reviews still produce fixed-length tensors, which is not only important for creating uniform batches but also becomes essential later when exporting the model to ONNX or TorchScript in Lesson 2 because both formats expect consistent input shapes.

After tokenization, we attach the corresponding labels by copying examples["label"] into the labels field. Hugging Face models expect this exact key during the forward pass when computing classification loss. This small remapping keeps the training step clean because the LightningModule can simply read batch["labels"] without any additional preprocessing or renaming logic.

The method finally returns a dictionary containing input_ids, attention_mask, and labels. Lightning and Hugging Face Datasets take care of converting these into PyTorch tensors later when we call .set_format("torch") in the setup() method. This tokenization pattern is highly scalable because it supports parallel CPU tokenization, remains deterministic for reproducibility, and works seamlessly in multi-GPU environments like DDP or FSDP. It is the same preprocessing approach used in large-scale NLP pipelines, from academic benchmarks to production systems handling millions of records.

Why This Approach Scales Well (MLOps Insight)

A major advantage of this design is how well it scales in real-world MLOps environments. Because _tokenize_function is used together with the Hugging Face .map() API, tokenization automatically runs in parallel across CPU workers. This gives you high throughput even when preprocessing millions of text samples. More importantly, the transformation is deterministic (the same input will always produce the same tokenized output), which is critical when you need reproducibility across training runs, machines, or distributed setups.

This pattern also integrates cleanly with multi-GPU training strategies (e.g., DDP and FSDP). Each GPU receives its own shard of the dataset and applies the exact same preprocessing logic, avoiding subtle mismatches between processes. Memory usage stays efficient as well, because batches are generated lazily and only the final tensorized output is kept in memory.

Finally, centralizing all preprocessing logic inside a single helper function makes the pipeline easy to maintain. Whether you’re training on a single GPU, scaling out to multiple nodes, or running the same preprocessing for offline batch inference, this structure remains robust. It’s the same approach used in production NLP systems where datasets can reach tens of millions of records and preprocessing needs to remain both fast and reliable.

Data Loaders: Batching Data for Training, Validation, and Testing

Once the dataset has been downloaded, tokenized, and formatted as tensors, the final component of the DataModule is producing PyTorch data loaders. These data loaders handle shuffling, batching, multiprocessing, and memory pinning, which are all needed to efficiently feed data into the model during training.

Here are the 3 data loader methods in our DataModule:

def train_dataloader(self):
    return DataLoader(
        self.train_dataset,
        batch_size=self.batch_size,
        shuffle=True,
        num_workers=self.num_workers,
        pin_memory=True
    )

def val_dataloader(self):
    return DataLoader(
        self.val_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        pin_memory=True
    )

def test_dataloader(self):
    return DataLoader(
        self.test_dataset,
        batch_size=self.batch_size,
        shuffle=False,
        num_workers=self.num_workers,
        pin_memory=True
    )

Let us walk through these line-by-line.

The final part of the DataModule defines the 3 data loaders Lightning will use for training, validation, and testing. Each method returns a standard PyTorch DataLoader, but Lightning wires them together automatically inside trainer.fit(), so you do not need to manually pass datasets around.

The train_dataloader() creates batches from the tokenized training set and enables shuffle=True. Shuffling is essential for training stability because it prevents the model from seeing samples in the same order every epoch, which helps it generalize better. The batch size and number of workers come directly from the DataModule’s configuration, allowing Hydra to control them externally.

The val_dataloader() follows the same structure, except it disables shuffling. Validation should always process data in a fixed order because we want deterministic metrics that reflect model performance, not random input order.

Finally, the test_dataloader() mirrors the validation setup (no shuffling and the same batching logic), ensuring that evaluation on the test set is consistent and repeatable.

All 3 data loaders also set pin_memory=True, a small but meaningful performance optimization. When using GPUs, pinned memory allows faster host-to-device transfers, reducing input bottlenecks during training. Lightning will automatically handle device placement of data, so the data loaders simply provide ready-to-consume batches.

Overall, these 3 methods complete the DataModule’s lifecycle: downloading, tokenizing, splitting, and packaging the data into efficient data loaders. With these in place, the Trainer can run end-to-end training without any additional data plumbing from your side, which provides exactly the modularity and cleanliness we want in an MLOps-focused codebase.

Why This Loader Setup Works Well for MLOps

This pattern brings several production advantages:

• Works seamlessly with DDP/FSDP across multiple GPUs: Lightning automatically replicates data loaders across processes
• Fully deterministic: Regenerating the same batches across runs supports experiment reproducibility
• Fast I/O due to multiprocessing + pinned memory: Critical for high-speed training workloads
• Clear boundaries for debugging: If a batch is wrong, you know the issue is in the DataModule, not the model

PyTorch LightningModule Explained: Building Modular Deep Learning Models

The LightningModule is the heart of the training pipeline.

While the DataModule organizes data, the LightningModule organizes learning (the model architecture, forward pass, loss computation, metrics, and optimization strategy).

In plain terms, the LightningModule is where “what the model is” and “how the model learns” are defined.

For sentiment classification, the module wraps a pretrained DistilBERT encoder, adds a small classification head, and implements the training/validation/test logic needed for stable and reproducible NLP experiments.

Initialization: Building the Model Architecture

class SentimentClassifier(pl.LightningModule):
    def __init__(
        self,
        model_name: str = "distilbert-base-uncased",
        num_labels: int = 2,
        dropout: float = 0.1,
        lr: float = 2e-5,
        weight_decay: float = 0.01,
        **kwargs
    ):
        super().__init__()
        self.save_hyperparameters()

        config = AutoConfig.from_pretrained(model_name)
        self.encoder = AutoModel.from_pretrained(model_name, config=config)

        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(config.hidden_size, num_labels)

        self.criterion = nn.CrossEntropyLoss()

        self.train_acc = Accuracy(task="multiclass", num_classes=num_labels)
        self.val_acc = Accuracy(task="multiclass", num_classes=num_labels)
        self.val_f1 = F1Score(task="multiclass", num_classes=num_labels, average="macro")

The constructor defines the core architecture and all hyperparameters needed for training. The use of self.save_hyperparameters() ensures that every setting (from model name to learning rate) is automatically captured in Lightning’s checkpoint files. This is extremely valuable in MLOps workflows because it makes every experiment self-describing and fully reproducible.

A DistilBERT encoder is loaded using the Hugging Face AutoModel, which instantly gives the model robust language understanding without any manual feature engineering. On top of that encoder, the model adds a simple dropout layer for regularization and a linear classifier to convert the encoded representation into sentiment logits. This small, clean architecture keeps training fast while still benefiting from the power of modern transformers.

Loss and metrics are also initialized here. Cross-entropy is the standard loss function for classification, while Accuracy and F1Score track performance during training and validation. Storing metrics inside the LightningModule keeps the entire training logic encapsulated cleanly in one place.

Forward Pass: How the Model Produces Predictions

def forward(self, input_ids, attention_mask):
    outputs = self.encoder(
        input_ids=input_ids,
        attention_mask=attention_mask
    )
    pooled = outputs.last_hidden_state[:, 0]
    logits = self.classifier(self.dropout(pooled))
    return logits

The forward method defines how a batch of tokenized text moves through the model. DistilBERT returns hidden states for all tokens, but for classification tasks it is sufficient to use the representation of the first token (the CLS token). The dropout + linear layer converts this pooled representation into raw class logits. Because we return logits directly, the LightningModule maintains full flexibility for computing loss and metrics in the training, validation, and test loops.

Training Step: Computing Loss and Logging Metrics

def training_step(self, batch, batch_idx):
    logits = self(batch["input_ids"], batch["attention_mask"])
    loss = self.criterion(logits, batch["labels"])
    preds = torch.argmax(logits, dim=1)
    acc = self.train_acc(preds, batch["labels"])

    self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True)
    self.log("train_acc", acc, on_step=True, on_epoch=True, prog_bar=True)
    return loss

The training_step defines exactly how the model learns from each batch. It runs a forward pass, computes the loss against the ground-truth labels, and updates accuracy. Lightning’s self.log handles aggregation across GPUs and ensures clean output in both the console and TensorBoard. Because Lightning abstracts away the optimizer step and gradient handling, this method focuses purely on the model’s learning logic rather than the boilerplate.

Validation Step: Tracking Generalization

def validation_step(self, batch, batch_idx):
    logits = self(batch["input_ids"], batch["attention_mask"])
    loss = self.criterion(logits, batch["labels"])

    preds = torch.argmax(logits, dim=1)
    acc = self.val_acc(preds, batch["labels"])
    f1 = self.val_f1(preds, batch["labels"])

    self.log("val_loss", loss, on_epoch=True, prog_bar=True)
    self.log("val_acc", acc, on_epoch=True, prog_bar=True)
    self.log("val_f1", f1, on_epoch=True)
    return loss

The validation step mirrors the training step but excludes gradient updates. It focuses entirely on monitoring how well the model generalizes. The addition of F1Score is especially useful in sentiment classification, where class imbalance is common and accuracy alone may hide poor performance on minority classes. Lightning automatically runs this method at the end of each epoch.

Test Step: Final Evaluation

def test_step(self, batch, batch_idx):
    logits = self(batch["input_ids"], batch["attention_mask"])
    loss = self.criterion(logits, batch["labels"])

    preds = torch.argmax(logits, dim=1)
    acc = Accuracy(task="multiclass", num_classes=self.num_labels).to(self.device)
    test_acc = acc(preds, batch["labels"])

    self.log("test_loss", loss, on_epoch=True)
    self.log("test_acc", test_acc, on_epoch=True)
    return loss

The test step follows the same pattern but uses dedicated metrics. This design keeps test evaluation separate from training and validation behaviors, which is useful when training and testing happen in different environments (e.g., offline batch scoring in production).

Optimizer Configuration: How the Model Learns

def configure_optimizers(self):
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in self.named_parameters()
                       if not any(nd in n for nd in no_decay)],
            "weight_decay": self.weight_decay,
        },
        {
            "params": [p for n, p in self.named_parameters()
                       if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]
    return torch.optim.AdamW(optimizer_grouped_parameters, lr=self.lr)

This method defines how parameters update during training. Transformer models benefit from AdamW, but they also require careful separation of weights that should and should not receive weight decay. Lightning lets you return the optimizer directly and handles everything else (e.g., multi-GPU synchronization, gradient scaling, and checkpointing). This keeps optimization logic concise while still following the best practices used in modern NLP training.

Why This Optimizer Pattern Matters in MLOps

This optimizer setup does more than just “train the model.”

By grouping parameters with and without weight decay, you ensure that every training run behaves consistently (regardless of hardware, environment, or number of GPUs). This is important for reproducibility because small optimizer differences can lead to diverging results in downstream evaluations.

Lightning also tracks the optimizer state inside checkpoints, which means that if training is resumed (perhaps in a new environment or after a failure), the run picks up with the exact same momentum buffers and learning dynamics. This is essential for building reliable training pipelines in real MLOps workflows.

Using Hydra Python Configuration Files for Reproducible ML Training

Hydra is the backbone of the training workflow in Lesson 1. Instead of hard-coding hyperparameters inside Python scripts, all settings for the model, data pipeline, and Trainer are stored cleanly inside YAML files. This keeps the codebase flexible, reproducible, and highly MLOps-friendly.

Hydra loads and composes these files automatically through the decorator in train.py, and you can override any value at runtime from the command line.

Let’s walk through each configuration file used in Lesson 1.

Root Configuration: configs/config.yaml

defaults:
  - model: distilbert
  - data: imdb
  - trainer: default
  - _self_

seed: 42

This is the entry point Hydra reads first.

The defaults block instructs Hydra to compose the final configuration from:

• configs/model/distilbert.yaml
• configs/data/imdb.yaml
• configs/trainer/default.yaml

The final line, seed: 42, becomes cfg.seed inside train.py, and controls reproducibility through pl.seed_everything().

This root file defines what experiment you’re running, while the nested configs define how each component behaves.

Model Configuration: configs/model/distilbert.yaml

model_name: distilbert-base-uncased
num_labels: 2
learning_rate: 2e-5
weight_decay: 0.01
dropout: 0.1

These values are passed directly into:

model = SentimentClassifier(**cfg.model)

Meaning:

• model_name: selects the Hugging Face backbone.
• num_labels: controls the classification head.
• learning_rate, weight_decay, and dropout: feed into the optimizer and architecture.

No values are hard-coded in Python, which is a major MLOps best practice that simplifies experiment tracking and hyperparameter sweeps.

Data Configuration: configs/data/imdb.yaml

batch_size: 8
max_length: 128
num_workers: 4

These values are injected directly into the DataModule:

datamodule = SentimentDataModule(
    model_name=cfg.model.model_name,
    **cfg.data
)

They control:

• how the data loader batches samples
• how long each sequence can be after tokenization
• how many CPU workers are used for preprocessing

This separation keeps the data pipeline reusable across models and experiments.

Trainer Configuration: configs/trainer/default.yaml

max_epochs: 3
devices: 1
accelerator: auto
precision: 32

These settings define how Lightning’s Trainer behaves:

trainer = pl.Trainer(**cfg.trainer, ...)

The values specify:

• number of training epochs
• how many devices (CPUs/GPUs) to use
• mixed precision strategy (FP32 for Lesson 1)
• automatic device selection (auto chooses CPU or GPU as available)

In Lesson 2, this same configuration pattern will allow you to enable DDP, FSDP, and mixed precision with zero Python code changes, requiring only new YAML overrides.

Why This Configuration System Is Critical for MLOps

Hydra gives you:

• clean separation of concerns
• reproducible experiments (config + code define a run)
• easy command-line overrides
• composable, hierarchical configurations
• effortless scaling for future lessons (distributed training, mixed precision, exports)

This structure is what makes the whole training pipeline “industrial-ready” instead of being a one-off script.

Building a PyTorch Lightning Training Pipeline with train.py

The training script is the glue that brings the entire pipeline together. Up to this point, you built a clean DataModule for handling data, a LightningModule for modeling logic, and a Hydra configuration system for reproducibility. Now, train.py uses all three to create a fully configurable training workflow without requiring manual training loops.

This script intentionally contains no data processing, no model logic, and no hard-coded hyperparameters. Instead, everything is driven by Hydra configs and Lightning abstractions, making the training pipeline reproducible, maintainable, and MLOps-friendly.

Let us walk through each component.

Hydra Configuration Loading

The script begins with Hydra’s decorator:

@hydra.main(config_path="../configs", config_name="config", version_base=None)
def main(cfg: DictConfig):

This single line gives the entire training pipeline a powerful configuration system. Hydra automatically:

• loads all YAML files under configs/
• composes them into a single cfg object
• injects trainer, data, and model settings dynamically
• allows overrides at runtime:

python src/train.py trainer.max_epochs=5 model.lr=3e-5

This configuration-first approach is a hallmark of modern MLOps workflows: the training script stays stable while configs change between runs.

Ensuring Reproducibility

Before building any components, you fix the training seed:

pl.seed_everything(cfg.seed, workers=True)

This sets random seeds across PyTorch, CUDA, NumPy, and Python itself. With workers=True, Lightning synchronizes DataLoader workers as well.

In MLOps pipelines, deterministic behavior is essential for:

• debugging
• regression testing
• comparing experiments fairly
• tracking drift in downstream metrics

Seeding makes every run repeatable.

Initializing the DataModule and Model

Next, the script constructs the full data pipeline:

datamodule = SentimentDataModule(
    model_name=cfg.model.model_name,
    **cfg.data
)

Hydra provides the dataset name, batch size, sequence length, tokenizer model, and num_workers.

Because all preprocessing logic is encapsulated inside the DataModule, the training script remains clean and focused.

Then the model is initialized:

model = SentimentClassifier(**cfg.model)

Every hyperparameter (e.g., learning rate, dropout, and number of labels) flows directly from the configuration. This separation between code and configuration makes the training pipeline:

• easier to maintain
• safer to scale
• more reproducible across teams and environments

Configuring Callbacks (Checkpointing + LR Monitoring)

Callbacks enable essential training behaviors without cluttering the model code.

ModelCheckpoint

checkpoint_callback = ModelCheckpoint(
    dirpath="checkpoints",
    filename="best-{epoch:02d}-{val_acc:.4f}",
    monitor="val_acc",
    mode="max",
    save_top_k=1,
    save_last=True,
    verbose=True
)

This callback:

• saves the best model based on validation accuracy
• keeps the last checkpoint for safety
• stores everything in a structured directory (checkpoints/)
• generates readable filenames

This makes it trivial to reload the best model for evaluation or deployment.

LearningRateMonitor

lr_monitor = LearningRateMonitor(logging_interval="step")

This tracks how the learning rate changes during training. LR patterns are important indicators of:

• instability
• vanishing gradients
• misconfigured schedulers
• early plateauing

Both callbacks are added to the Trainer:

callbacks = [checkpoint_callback, lr_monitor]

Logging with TensorBoard

Lightning integrates cleanly with TensorBoard through:

logger = TensorBoardLogger(
    save_dir="logs",
    name="sentiment_classifier"
)

The logger records:

• loss curves
• accuracy
• learning rate
• hyperparameters (optional)

You can visualize results with:

tensorboard --logdir logs/

In Lesson 2 and later modules, this logging layer leads naturally into MLflow, W&B, or Langfuse for real MLOps observability.

Constructing the Trainer

Lightning’s Trainer consolidates all runtime behavior into one object:

trainer = pl.Trainer(
    **cfg.trainer,
    callbacks=callbacks,
    logger=logger,
    enable_progress_bar=True
)

Hydra injects:

• accelerator (cpu/gpu)
• devices
• precision
• max epochs
• gradient clipping
• logging frequency

This makes the script hardware-agnostic. In Lesson 2, you will enable DDP, FSDP, and mixed precision simply by changing the trainer config without modifying the Python code.

Launching Training

Finally, the entire pipeline is executed with:

trainer.fit(model, datamodule)

This one method call triggers:

• dataset downloading (prepare_data)
• tokenization (setup)
• data loader creation
• training and validation loops
• checkpointing
• logging

After training completes, the script prints useful summary information:

best_checkpoint = checkpoint_callback.best_model_path
val_acc = trainer.callback_metrics.get("val_acc", 0)
val_loss = trainer.callback_metrics.get("val_loss", 0)

This delivers the key metrics needed for evaluation or comparison with future runs.

Running Model Inference with PyTorch Lightning and DistilBERT

Once training is complete, the next step in any ML or MLOps workflow is deploying the model for inference. Lesson 1 focuses on offline inference (i.e., running predictions locally using the saved checkpoints). Your inference script is flexible and production-ready, supporting 4 different usage modes:

• Single-text prediction
• Batch prediction from a file
• Interactive command-line mode
• Demo mode with predefined reviews

This section walks through every component of the script and explains how it all fits together.

Here is the full header and imports:

import argparse
from pathlib import Path
import torch
from transformers import AutoTokenizer
from model_module import SentimentClassifier

The imports bring in argparse for command-line parsing, torch for tensor and device management, AutoTokenizer for text preprocessing, and your trained LightningModule class SentimentClassifier. Path is imported but not strictly required in this version; that’s fine and easy to reuse later.

Loading the Trained Model

def load_model(checkpoint_path: str):
    """Load trained model from checkpoint."""
    print(f"Loading model from {checkpoint_path}...")
    # Load to CPU first to avoid device mismatch issues
    # Works for checkpoints saved on any device (CPU/CUDA/MPS/multi-GPU)
    model = SentimentClassifier.load_from_checkpoint(
        checkpoint_path,
        map_location="cpu"
    )
    model.eval()
    model.freeze()
    print("✅ Model loaded successfully!\n")
    return model

load_model() encapsulates everything needed to restore the model from disk. It prints which checkpoint is being loaded for transparency, then uses Lightning’s load_from_checkpoint() to rebuild the full SentimentClassifier with weights and hyperparameters.

Setting map_location="cpu" guarantees that the checkpoint can be loaded no matter where it was trained (single GPU, multi-GPU, CPU, MPS). Calling eval() puts the model into inference mode, turning off dropout and other training-only behaviors, and freeze() disables gradients so the model becomes a pure forward-pass computation graph. Finally, it prints a success message and returns the ready-to-use model.

Predicting Sentiment for a Single Text

def predict_sentiment(model, tokenizer, text: str, device: str = "auto"):
    """Predict sentiment for a single text."""
    # Tokenize input
    encoding = tokenizer(
        text,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    
    # Move to device
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    
    input_ids = encoding["input_ids"].to(device)
    attention_mask = encoding["attention_mask"].to(device)
    model = model.to(device)
    
    # Get prediction
    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        probs = torch.softmax(logits, dim=1)
        pred_class = torch.argmax(probs, dim=1).item()
        confidence = probs[0, pred_class].item()
    
    sentiment = "Positive 😊" if pred_class == 1 else "Negative 😞"
    
    return {
        "text": text,
        "sentiment": sentiment,
        "confidence": confidence,
        "label": pred_class
    }

predict_sentiment() is the core prediction helper. It starts by tokenizing the input text using the same tokenizer settings as training: max_length=128, truncation=True, and padding="max_length". This ensures the input tensor shape is fixed and consistent with what DistilBERT expects, which is also helpful later when exporting to ONNX/TorchScript.

The device argument defaults to "auto", which makes the function portable. If "auto" is requested, it checks whether CUDA is available and chooses GPU when possible, otherwise falls back to CPU. Both the input_ids and attention_masktensors are moved onto that device, and the model itself is also transferred using model.to(device). This guarantees that all tensors and the model live on the same device, avoiding runtime errors.

Inside a torch.no_grad() block, the function performs a forward pass: it feeds input_ids and attention_mask into the model, gets raw logits, converts them into probabilities with torch.softmax, and then finds the predicted class with torch.argmax. The .item() calls convert these tiny tensors into regular Python scalars. The sentiment string maps class 1 to positive and everything else to negative, adding emoji for an immediately readable result. Finally, the function returns a dictionary containing the original text, the sentiment label, confidence score, and numeric class. This output is easy to print, log, or send through an API.

Interactive Inference Mode

def interactive_mode(model, tokenizer):
    """Interactive inference mode."""
    print("=" * 70)
    print("Interactive Sentiment Analysis")
    print("=" * 70)
    print("\nType your movie review and press Enter.")
    print("Type 'quit' or 'exit' to stop.\n")
    
    while True:
        text = input(" Review: ").strip()
        
        if text.lower() in ['quit', 'exit', 'q']:
            print("\n Goodbye!")
            break
        
        if not text:
            continue
        
        result = predict_sentiment(model, tokenizer, text)
        
        print(f"\n{'='*70}")
        print(f"Sentiment: {result['sentiment']}")
        print(f"Confidence: {result['confidence']:.2%}")
        print(f"{'='*70}\n")

interactive_mode() turns the script into a small command-line app. It prints a header and basic instructions, then enters an infinite loop. Each iteration reads a line from input(), trims whitespace, and checks for exit commands (quit, exit, q). Empty lines are ignored so the user doesn’t accidentally trigger computation.

For valid text, it calls predict_sentiment() with the provided model and tokenizer. The results are printed in a nicely formatted block, showing the sentiment label and confidence percentage. This mode is perfect for quickly sanity-checking the model, exploring its behavior, and giving students an immediate feel for how the classifier responds to different reviews.

Batch Inference from a Text File

def batch_inference(model, tokenizer, input_file: str):
    """Run inference on multiple texts from file."""
    print(f"Reading texts from {input_file}...")
    
    with open(input_file, 'r') as f:
        texts = [line.strip() for line in f if line.strip()]
    
    print(f"Running inference on {len(texts)} samples...\n")
    print("=" * 70)
    
    results = []
    for i, text in enumerate(texts, 1):
        result = predict_sentiment(model, tokenizer, text)
        results.append(result)
        
        print(f"\n{i}. {text[:60]}{'...' if len(text) > 60 else ''}")
        print(f"   → {result['sentiment']} ({result['confidence']:.2%})")
    
    print("\n" + "=" * 70)
    print(f"\n Processed {len(results)} samples")
    
    # Summary statistics
    positive = sum(1 for r in results if r['label'] == 1)
    negative = len(results) - positive
    avg_confidence = sum(r['confidence'] for r in results) / len(results)
    
    print(f"\n Summary:")
    print(f"   Positive: {positive} ({positive/len(results):.1%})")
    print(f"   Negative: {negative} ({negative/len(results):.1%})")
    print(f"   Avg Confidence: {avg_confidence:.2%}")

batch_inference() is designed for real-world workflows where you want to score multiple reviews at once. It accepts a path to a file, reads all non-empty lines into a list, and reports how many samples it found. For each line, it calls predict_sentiment() and prints out a truncated version of the text plus its predicted sentiment and confidence. Results are collected into a list so you can compute summary statistics afterward.

At the end, it computes how many predictions were positive vs negative and the average confidence across all samples. Those aggregate metrics are useful when running offline evaluations, batch scoring jobs, or quick experiments. This pattern also generalizes well to production pipelines where you might later replace the file input with a database, message queue, or data warehouse.

Demo Mode with Sample Reviews

def demo_mode(model, tokenizer):
    """Run demo with sample reviews."""
    print("=" * 70)
    print("Demo: Sample Movie Reviews")
    print("=" * 70)
    
    sample_reviews = [
        "This movie was absolutely fantastic! Best film I've seen this year.",
        "Terrible movie, complete waste of time. Boring and poorly acted.",
        "The plot was confusing and the ending made no sense."
    ]
    
    print(f"\nTesting {len(sample_reviews)} sample reviews...\n")
    
    for i, text in enumerate(sample_reviews, 1):
        result = predict_sentiment(model, tokenizer, text)
        print(f"{i}. {text}")
        print(f"   → {result['sentiment']} ({result['confidence']:.2%})\n")
    
    print("=" * 70)

demo_mode() is a curated, zero-configuration way to showcase the model. It defines a fixed list of sample movie reviews that cover positive, negative, and neutral-ish cases. It prints a header, shows how many samples will be tested, and then iterates through each review, calling predict_sentiment() under the hood. The output is a neat numbered list with the original sentence and the corresponding sentiment and confidence.

This is ideal for teaching, demos, and quick regression checks after changing code or dependencies. If the demo suddenly behaves strangely, you know something is wrong.

Command-Line Interface and Mode Routing

def main():
    parser = argparse.ArgumentParser(description="Sentiment Classification Inference")
    parser.add_argument(
        "--checkpoint",
        type=str,
        required=True,
        help="Path to model checkpoint (e.g., checkpoints/best-*.ckpt)"
    )
    parser.add_argument(
        "--text",
        type=str,
        help="Single text to classify"
    )
    parser.add_argument(
        "--input-file",
        type=str,
        help="File with texts to classify (one per line)"
    )
    parser.add_argument(
        "--demo",
        action="store_true",
        help="Run demo with sample reviews"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="auto",
        choices=["auto", "cpu", "cuda", "mps"],
        help="Device to run inference on"
    )
    
    args = parser.parse_args()
    
    # Load model
    model = load_model(args.checkpoint)
    
    # Load tokenizer
    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model.hparams.model_name)
    print("Tokenizer loaded\n")
    
    # Run appropriate mode
    if args.demo:
        demo_mode(model, tokenizer)
    elif args.text:
        result = predict_sentiment(model, tokenizer, args.text, args.device)
        print("=" * 70)
        print(f"Text: {result['text']}")
        print(f"Sentiment: {result['sentiment']}")
        print(f"Confidence: {result['confidence']:.2%}")
        print("=" * 70)
    elif args.input_file:
        batch_inference(model, tokenizer, args.input_file)
    else:
        interactive_mode(model, tokenizer)


if __name__ == "__main__":
    main()

The main() function wires everything into a clean CLI interface. argparse defines all supported flags:

• --checkpoint (required): tells the script which .ckpt file to load
• --text: enables single-text prediction
• --input-file: enables batch inference from a file
• --demo: triggers the curated demo mode
• --device: lets you override the device choice (auto, cpu, cuda, mps)

After parsing arguments, main() calls load_model() to restore the trained Lightning model. It then loads the tokenizer using AutoTokenizer.from_pretrained(model.hparams.model_name), which reads the model name stored in the checkpoint hyperparameters. This ensures inference uses the exact same tokenizer that training used.

Finally, it chooses which mode to run based on the arguments:

• If --demo is set: run demo_mode()
• Else if --text is provided: run predict_sentiment() once and print the result
• Else if --input-file is provided: run batch_inference()
• Else: fall back to interactive_mode()

The if __name__ == "__main__": guard makes this file executable as a script and keeps it import-safe if you ever reuse these helpers in a larger app or API.

Training and Evaluating a DistilBERT Model with PyTorch Lightning and Hydra

Now that we have built our DataModule, LightningModule, and training script, it is time to run everything end-to-end.

This section walks you through:

• launching training with Hydra
• visualizing metrics in TensorBoard
• running inference in multiple modes
• understanding how Lightning automatically selects the best available device (CUDA, MPS, or CPU)

To learn how to run and evaluate your sentiment classifier, just keep reading.

Running Training from the Command Line

You can start a basic training run using:

$ python src/train.py

Lightning will automatically download the IMDB dataset, initialize DistilBERT, set seeds, and begin training.

$ python src/train.py
================================================================================
LESSON 1: PyTorch Lightning Fundamentals with Hydra
================================================================================
📋 Configuration:
model:
  model_name: distilbert-base-uncased
  num_labels: 2
  dropout: 0.1
  lr: 2.0e-05
  weight_decay: 0.01
data:
  dataset_name: imdb
  max_length: 128
  batch_size: 8
  num_workers: 4
trainer:
  max_epochs: 3
  accelerator: auto
  devices: 1
  precision: 32
  log_every_n_steps: 10
  check_val_every_n_epoch: 1
  deterministic: false
seed: 42
Seed set to 42
🌱 Seed set to: 42
   → Ensures reproducible results across runs
📊 Initializing DataModule...
   Dataset: imdb
   Batch size: 8
   Max sequence length: 128
🤖 Initializing Model...
   Model: distilbert-base-uncased
   Learning rate: 2e-05
   Number of labels: 2
config.json: 100%|██████████████████████████████| 483/483 [00:00<00:00, 682kB/s]
model.safetensors:   0%|                             | 0.00/268M [00:00<?, ?B/s]

This screenshot typically includes:

• Hydra-composed config
• “Seed set to …”
• DataModule initialization logs
• Model initialization
• Trainer configuration summary

It gives readers immediate confidence that the training pipeline is well-structured and reproducible.

Overriding Hyperparameters with Hydra

One of Hydra’s biggest strengths is that it allows you to change any configuration value without modifying a single line of code.

For example, you can increase the number of training epochs:

$ python src/train.py trainer.max_epochs=5

Or explore a different learning rate:

$ python src/train.py model.lr=3e-5

Or change the batch size:

$ python src/train.py data.batch_size=16

Hydra composes the full configuration at runtime and applies the CLI override (data.batch_size=16). The screenshot shows the merged YAML, seed setup, DataModule initialization, model initialization, and callback configuration, all without changing a single line of code. This illustrates the core MLOps benefit of configuration-driven training.

This is a major MLOps advantage because configuration-driven training allows experiments to be repeated, compared, and automated easily.

Viewing Logs and Metrics in TensorBoard

All training metrics are logged to:

logs/sentiment_classifier/

You can launch TensorBoard using:

$ tensorboard --logdir logs/

This screenshot shows how TensorBoard automatically tracks metrics (e.g., val_acc and val_loss) for each Hydra-versioned training run. Because the entire pipeline is configuration-driven, you can compare experiments side-by-side by smoothing curves, inspecting run histories, and validating that model performance trends remain consistent across seeds or hyperparameter overrides.

A typical figure here would show:

• Training vs. validation loss
• Validation accuracy
• Learning rate
• Iteration timeline

This screenshot reinforces that Lightning + TensorBoard gives you well-structured experiment tracking out of the box.

Where Checkpoints Are Saved

Lightning automatically saves your best checkpoint based on validation accuracy:

checkpoints/best-epoch=XX-val_acc=YY.ckpt

It also saves a last.ckpt for safety.

Lightning’s ModelCheckpoint callback saves the best-performing model (based on val_acc) and the final “last.ckpt” snapshot for safety. Because training is fully configuration-driven, each Hydra-run creates its own versioned checkpoint set, making it easy to resume, compare, or deploy models in a reproducible MLOps workflow.

This is often the moment readers see how cleanly Lightning organizes artifacts, which ties directly into the next DVC module.

Running Inference (Multiple Modes)

Once training completes, you can use the inference script in 3 different modes.

A) Single-Text Prediction

$ python src/inference.py --checkpoint "checkpoints/best-epoch=00-val_acc=0.8797.ckpt" --text "This movie was amazing!"

The screenshot shows the model being loaded from the selected checkpoint, the tokenizer initialization, and the final prediction with confidence for the given input text.

Lightning auto-loads the checkpoint, loads the tokenizer, moves tensors to the correct device, and returns a structured prediction.

B) Batch Inference from File

Create a file such as the following:

samples.txt
This movie was incredible.
The film was a disappointment.
I loved every character!

Then run:

$ python src/inference.py \
    --checkpoint checkpoints/best-epoch=*.ckpt \
    --input-file samples.txt

The screenshot shows the model loading, tokenizer initialization, per-review predictions with confidence scores, and a final summary of positive and negative counts. This demonstrates how the inference pipeline scales cleanly from single inputs to larger batches, which is a key requirement in practical MLOps workflows.

The script also prints a small summary (positive %, negative %, average confidence).

C) Interactive Mode

$ python src/inference.py \
    --checkpoint checkpoints/best-epoch=*.ckpt

You can then type reviews manually.

Figure 6 shows the interactive prompt, a sample user review (“The movie was good — just not great.”), and the model’s corresponding prediction with confidence. This mode is ideal for quick testing and showcasing the end-to-end inference flow after training your Lightning model.

This mode is great for demos, lightweight testing, or API prototyping.

Device Auto-Selection (CPU, GPU, and MPS)

Your inference script includes this logic:

if device == "auto":
    device = "cuda" if torch.cuda.is_available() else "cpu"

When device is set to "auto", the script selects a CUDA-enabled graphics processing unit (GPU) if one is available. Otherwise, it selects the central processing unit (CPU). On Macs with Apple Silicon, you can also choose Metal Performance Shaders (MPS) explicitly by setting the device to "mps".

Lightning also handles device selection on the training side:

• If a GPU is available and the trainer is configured to use it, Lightning uses the GPU.
• If multiple GPUs are configured, Lightning uses DDP (Lesson 2).
• If no supported GPU is available, Lightning falls back to the CPU.

Figure 7 confirms that Lightning detected an MPS-capable GPU on macOS and will use it for training. It also reports that no Tensor Processing Unit (TPU) cores are available. This automation allows the training code to run across different hardware configurations without modification.

This flexibility dramatically improves developer experience (DX) because your code can run on any machine.

Summary: Why This Section Matters for MLOps

This “Run and Evaluate” section ties the entire workflow together:

• Hydra makes experiments repeatable and configurable.
• Lightning makes training structured and supports deterministic execution when configured appropriately.
• TensorBoard provides observability.
• Checkpoints preserve model state for downstream deployment.
• The inference script demonstrates end-to-end functionality.

This is precisely the level of workflow maturity expected in modern MLOps and LLMOps environments.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: August 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you learned how PyTorch Lightning and Hydra work together to bring structure, reproducibility, and MLOps-friendly practices to a modern NLP training pipeline. Instead of writing long, error-prone training loops, Lightning allowed you to organize your project into clean, isolated components: the LightningDataModule for handling all data operations and the LightningModule for model logic, optimization, and metrics. Hydra completed the picture by giving you a powerful configuration system that let you change hyperparameters, batch sizes, model choices, or even hardware strategies from the command line without touching a single line of Python.

You trained a sentiment classifier on the IMDB dataset using a fully modular workflow, ran experiments with different overrides, visualized metrics in TensorBoard, and saved reproducible checkpoints automatically. Most importantly, you saw how clean architecture and configuration-driven design make your training pipeline easier to extend, debug, monitor, and deploy. These capabilities are exactly what real-world MLOps demands.

In the next lesson, we will take this foundation and push it into production territory. You will learn how to scale the same codebase with distributed training (DDP and FSDP), mixed precision, and model export to ONNX and TorchScript, all while keeping the project structure identical. If you want to turn this into a robust, deployment-ready training workflow, just keep reading.

Citation Information

Singh, V. “Training with PyTorch Lightning: Structured MLOps Development,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/5fe4l

@incollection{Singh_2026_training-w-pytorch-lightning-structured-mlops-development,
  author = {Vikram Singh},
  title = {{Training with PyTorch Lightning: Structured MLOps Development}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/5fe4l},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post Training with PyTorch Lightning: Structured MLOps Development appeared first on PyImageSearch.

Running Gemma 4 Locally: Ollama, llama.cpp, MLX, and More

In the first part of this series, we explored the architecture behind Gemma 4, including its multimodal design, long-context capabilities, Mixture-of-Experts routing, Per-Layer Embeddings, and the engineering decisions that allow the model family to scale from lightweight on-device deployments to large production systems. We also examined the different Gemma 4 variants, their hardware requirements, and the reasoning capabilities that make the model family one of the most compelling open-weight releases available today.

In the second part, we moved from theory to practice. Using Hugging Face Transformers, we built multimodal applications capable of processing images, videos, audio, and structured outputs. Along the way, we saw how Gemma 4 handles tasks such as screenshot-to-code generation, object detection, image captioning, multimodal function calling, and audio understanding from a single unified interface.

But there is still an important question left unanswered:

How do we actually run Gemma 4 outside of a Python notebook?

For many developers, deploying and interacting with a model is just as important as understanding its architecture or capabilities. A model that performs well inside a notebook is useful, but a model that can run locally on a laptop, power desktop applications, expose APIs, or operate entirely offline opens up an entirely different set of possibilities.

Fortunately, Gemma 4 has quickly gained support across the local AI ecosystem. Whether you prefer a simple one-command setup, a graphical desktop interface, a highly optimized C++ inference engine, a runtime specifically designed for Apple Silicon, or even JavaScript-based local inference, there is already a mature deployment path available.

In this lesson, we will explore five of the most popular ways to run Gemma 4 locally:

• Ollama: fast setup and developer-friendly APIs
• LM Studio: a graphical desktop experience
• llama.cpp: efficient GGUF (GPT-Generated Unified Format)-based inference and maximum control
• MLX: optimized performance on Apple Silicon devices
• Transformers.js: running Gemma 4 directly from Node.js using ONNX models

By the end of this guide, you will know how to download Gemma 4, perform local text and multimodal inference, expose OpenAI-compatible APIs, and choose the deployment runtime that best fits your workflow.

In the next lesson, we will build upon the Transformers.js foundation introduced here and take things one step further by running Gemma 4 entirely inside the browser, enabling fully client-side AI applications without requiring a Python backend or native inference runtime.

This lesson is the 3rd in a 5-part series on Google DeepMind’s Gemma 4:

1. Google DeepMind’s Gemma 4: MoE, Efficiency Tricks, and Benchmarks
2. Building Multimodal AI Applications with Gemma 4 and Transformers
3. Running Gemma 4 Locally: Ollama, llama.cpp, MLX, and More (this tutorial)
4. Lesson 4
5. Lesson 5

To learn how to run Gemma 4 entirely on your own hardware, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Running Gemma 4 with Ollama for Local AI Inference

Among all the local inference solutions available today, Ollama is arguably the fastest and easiest way to get started with Gemma 4. With just a few terminal commands, it can download the model, manage different versions, launch a local inference server, and expose an OpenAI-compatible API (Application Programming Interface), all without requiring any manual configuration.

Instead of downloading model weights, setting up inference backends, or configuring server endpoints yourself, Ollama automates the entire deployment process behind a simple command-line interface. This makes it an excellent choice for developers who want to start experimenting with Gemma 4 immediately while still having access to production-friendly APIs.

In this section, we will install Ollama, download the Gemma 4 model, perform both text and vision inference from the terminal, and interact with the model through its built-in REST API.

Installing Ollama

The first step is installing Ollama on your machine.

curl -fsSL https://ollama.com/install.sh | sh

This command downloads the official installation script directly from the Ollama website and executes it locally.

After the installation completes, we can verify that Ollama was installed successfully:

ollama --version

If the command prints a version number, the installation is working correctly and we are ready to download Gemma 4.

Downloading the Gemma 4 Model

Ollama stores models locally and downloads them on demand.

To pull the Gemma 4 E2B model, run:

ollama pull gemma4:e2b

The pull command downloads the model weights and stores them inside Ollama’s local model registry.

Here:

• gemma4: specifies the model family
• e2b: refers to the Gemma 4 E2B variant

Once the download finishes, the model becomes available for local inference without requiring any additional configuration.

Running Gemma 4 Interactively

With the model downloaded, we can launch an interactive chat session:

ollama run gemma4:e2b

This command starts Gemma 4 and opens a terminal-based conversation interface.

You can now enter prompts directly into the terminal and receive responses from the model in real time.

>>> explain the mixture-of-experts architecture.
Thinking...
Here's a thinking process that leads to the suggested explanation:
1.  **Deconstruct the Request:** The user wants an explanation of the "Mixture-of-Experts (MoE) architecture."
2.  **Determine the Core Context (Where does MoE live?):** MoE is primarily a model architecture, most famously applied to Large Language Models (LLMs) like those from Google (GLaM/PaLM/Gemini) 
and Meta.
3.  **Establish the Problem MoE Solves (Why MoE?):** Standard dense models (like a traditional Transformer) are computationally expensive. They require massive parameter counts and massive 
computations for *every* input, even for simple tasks.
    *   *Goal:* Make models larger (more capable) without increasing the computational cost per inference/training step proportionally.
….
….

To exit the session, type:

/bye

or press:

Ctrl + D

This interactive mode is useful when experimenting with prompts, testing reasoning capabilities, or quickly validating model behavior before integrating it into an application.

Image Understanding with Gemma 4

Gemma 4 is a multimodal model, meaning it can process both text and images.

Ollama exposes this capability through a simple command:

ollama run gemma4:e2b "caption this image /Users/cosmo3769/Desktop/venice.jpg"

In this example, Gemma 4 receives an image and generates a textual description of its contents.

Behind the scenes, Ollama loads the image, converts it into visual embeddings using Gemma 4’s vision encoder, and passes those embeddings into the language model alongside the text prompt.

From the user’s perspective, image understanding works exactly like a standard text prompt.

Added image '/Users/cosmo3769/Desktop/venice.jpg'
Thinking...
Here's a thinking process for generating the caption:
1.  **Analyze the Image:**
    *   **Subject:** A large, classical-style building (a temple/fascia) situated over water.
    *   **Setting:** A canal/river scene. There are wooden pilings in the water, suggesting a waterfront or mooring area.
    *   **Architecture:** The building is ornate, stone, and features columns and pediments. It looks like a palazzo or a significant public structure.
    *   **Context/Location Clues (Inferred):** The architecture style strongly suggests Italy, likely Venice or a similar canal city.
    *   **Foreground/Background:** Wooden waterboats/barges are visible. There are old European buildings lining the far bank. A bridge/structure is visible on the right. The light suggests a 
sunny day.
2.  **Identify Key Themes & Mood:**
    *   History/Antiquity
    *   Waterways/Canals
    *   European Charm
    *   Architecture/Grandeur
3.  **Determine the Goal:** Create a compelling caption for this scenic, historical photo.
4.  **Draft Caption Options (Categorized):**
    *   *Descriptive/Literal:* Focus on what is seen.
    *   *Evocative/Poetic:* Focus on the feeling and atmosphere.
    *   *Location-Specific (If known):* If we assume Venice/Italy.
    *   *Travel/Wanderlust Focused:* Encouraging travel.
5.  **Refine and Select the Best Options (Adding specific details):**
    *   *Focus on the contrast:* The grandeur of the building vs. the water.
    *   *Focus on the atmosphere:* Serene, timeless.
    *   *Focus on the setting:* Venetian/Italian canals.
6.  **Final Polish (Generating a variety of styles):** (This leads to the suggested captions below.)
...done thinking.
Here are several options for captioning this beautiful image, depending on the tone you want to convey:
**1. Descriptive & Historical:**
*   "A glimpse of timeless elegance and history along the canals. The grandeur of the architecture reflected in the calm waters."
*   "Ancient beauty meets the waterways. Exploring the historic facades and hidden corners of this European city."
*   "The stunning architecture overlooking the tranquil waters. A true masterpiece of history."
**2. Evocative & Poetic:**
*   "Where history flows on water. Serenity and splendor in every frame."
*   "Lost in the charm of the canals, where old stones whisper tales of the past."
*   "Golden light on ancient waters. A perfect moment of European serenity."
**3. Travel & Wanderlust Focused:**
*   "Dreaming of Italian canals and timeless beauty. 🇮🇹"
*   "Wanderlust activated! Finding magic in the waterways of Europe."
*   "Exploring the hidden gems of a historic city by the water."
**4. Short & Punchy (Good for Instagram):**
*   "Canal views and classical charm."
*   "Waterways of wonder."
*   "Historic beauty afloat."
*   "Venetian dreams."
***
**✨ Pro-Tip: Add Relevant Hashtags**
#Venice #Italy #Canals #History #Architecture #Travel #Europe #HistoricCity #WaterViews #TravelGram

Using the OpenAI-Compatible API

One of Ollama’s most useful features is its OpenAI-compatible API.

When Ollama is running, it automatically exposes an HTTP endpoint that mimics the OpenAI Chat Completions API. This means existing applications built for OpenAI often require only a single URL change to work with local models.

We can send a request using cURL:

curl http://localhost:11434/v1/chat/completions \
 -H "Content-Type: application/json" \
 -d '{
   "model": "gemma4:e2b",
   "messages": [
     {"role": "user", "content": "Explain how transformers work in two sentences."}
   ]
 }'

Let us break down the request:

• model: specifies which local model should handle the request
• messages: contains the conversation history
• role: identifies whether a message comes from the user, assistant, or system
• content: contains the actual prompt

The response format mirrors OpenAI’s API structure, making it straightforward to integrate Ollama into existing AI applications.

Sending Images via the Ollama API

Ollama also supports multimodal requests through its native API.

The following example sends an image directly to Gemma 4:

curl http://localhost:11434/api/generate -d "{
 \"model\": \"gemma4:e2b\",
 \"prompt\": \"Describe what you see in this image\",
 \"images\": [\"$(base64 -i /Users/cosmo3769/Desktop/venice.jpg)\"],
 \"stream\": false
}"

Unlike the OpenAI-compatible endpoint, this API accepts image data directly.

Notice the use of:

base64 -i image.jpg

Images cannot be transmitted directly inside JSON payloads, so we first encode the image as a Base64 string. Ollama then decodes the image, processes it through Gemma 4’s vision encoder, and generates a textual response describing what it sees.

The "stream": false parameter tells Ollama to return the complete response in a single payload instead of streaming tokens incrementally.

At this point, we have already covered the majority of day-to-day Gemma 4 workflows:

• Interactive chat
• Text generation
• Multimodal image understanding
• OpenAI-compatible APIs
• Programmatic image inference

For many developers, Ollama is all that is needed to begin building local Gemma 4 applications.

To learn more about Ollama, you can refer to this blog from PyImageSearch.

Running Gemma 4 with llama.cpp and GGUF Models

While Ollama prioritizes simplicity and ease of use, llama.cpp is built for performance, flexibility, and fine-grained control. It is one of the most popular inference engines for running large language models locally, offering efficient CPU execution, hardware acceleration, and native support for GGUF models.

Unlike Ollama, which abstracts away much of the deployment process, llama.cpp gives us direct control over every aspect of inference. We explicitly choose the model file, configure hardware acceleration, launch inference servers, and customize runtime parameters to suit our hardware and workload. This level of control makes llama.cpp an excellent choice for developers who want to optimize performance or integrate local models into custom applications.

In this section, we will build llama.cpp from source, download a Gemma 4 GGUF model, run inference from the command line, launch an OpenAI-compatible API server, and finally enable multimodal image understanding using Gemma 4’s MMProj (multimodal projector) model.

Installing Build Dependencies

Before building llama.cpp, we first need to install the required development tools.

xcode-select --install
clang --version

The first command installs Apple’s Command Line Tools, which include the Clang compiler and other utilities required for compiling C++ projects.

After the installation completes, we can verify that Clang is available by checking its version.

brew install git cmake
git --version
cmake --version

Next, install Git and CMake using Homebrew.

Here:

• Git: is used to clone the llama.cpp source code.
• CMake: generates platform-specific build files and manages the compilation process.

Running the version commands confirms that both tools were installed successfully.

Cloning the Repository

git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp

Once the required tools are installed, we can download the llama.cpp source code directly from GitHub.

The git clone command downloads the latest version of the repository to your local machine.

We then change into the project directory, where all remaining build and inference commands will be executed.

Building llama.cpp

Next, we compile llama.cpp from source.

mkdir build
cd build
cmake ..

Let us briefly examine these commands.

First, we create a separate build directory. Keeping compiled files separate from the source code makes it easier to rebuild or clean the project later.

The command: cmake .. configures the project and generates the appropriate build files for your system.

After configuration completes, we compile the project using:

cmake --build . --config Release

The Release configuration enables compiler optimizations, resulting in faster inference performance than a debug build.

Once the compilation finishes, the generated executables can be found inside the build/bin directory.

This directory contains utilities such as:

• llama-cli
• llama-server
• model conversion tools
• quantization utilities

These binaries form the core of the llama.cpp inference ecosystem.

Enabling Apple Metal Acceleration

If you are using an Apple Silicon Mac, llama.cpp can offload model computation to the GPU through Apple’s Metal framework.

To enable Metal support, rebuild llama.cpp with the following configuration:

cmake -B build -DGGML_METAL=ON
cmake --build build --config Release

The GGML_METAL=ON flag enables GPU acceleration during compilation.

Compared to CPU-only inference, Metal acceleration can significantly improve generation speed, particularly for larger Gemma 4 models.

Verifying the Installation

Before downloading a model, it is a good idea to verify that the build completed successfully.

./build/bin/llama-cli --help

If the installation is successful, llama.cpp prints the available command-line options for llama-cli.

This confirms that the executable was built correctly and is ready to load Gemma 4.

Downloading and Locating a Gemma GGUF Model

Unlike Ollama, llama.cpp works directly with GGUF model files.

GGUF is a binary model format optimized for fast local inference and supports multiple quantization levels, allowing models to run efficiently on CPUs and consumer GPUs while reducing memory usage.

python3 -m venv .venv
source .venv/bin/activate

We first create a Python virtual environment:

./build/bin/llama-server -hf ggml-org/gemma-4-E2B-it-GGUF

Next, download the Gemma 4 GGUF model directly from Hugging Face:

The -hf option downloads the specified repository automatically and stores the model locally.

find ~ -name "*gemma*.gguf" 2>/dev/null

If you already have the model on your machine, you can locate it using:

This searches your home directory for all GGUF files containing “gemma” in their filename.

Running Gemma 4 from the Command Line

Once the model has been downloaded, we can perform inference directly from the terminal.

./build/bin/llama-cli \
 -m /Users/cosmo3769/.lmstudio/models/lmstudio-community/gemma-4-E2B-it-GGUF/gemma-4-E2B-it-Q4_K_M.gguf \
 -p "Explain the Mixture of Experts architecture" \
 -n 512 \
 -ngl 99

Here:

• -m: specifies the path to the GGUF model.
• -p: provides the prompt sent to Gemma 4.
• -n 512: allows the model to generate up to 512 new tokens.
• -ngl 99: offloads supported transformer layers to the GPU for faster inference. The exact value depends on your available GPU memory and hardware configuration.

NOTE: We are using the GGUF model from LM Studio which we downloaded earlier. We can also use it from hf but it will be the same.

After executing the command, Gemma 4 generates a response directly inside the terminal without requiring a separate inference server.

Starting a Local API Server

Instead of running one prompt at a time, we can expose Gemma 4 through a local HTTP server.

./build/bin/llama-server \
 -m /Users/cosmo3769/.lmstudio/models/lmstudio-community/gemma-4-E2B-it-GGUF/gemma-4-E2B-it-Q4_K_M.gguf \
 --host 0.0.0.0 \
 --port 8080 \
 -n 512

Here:

• --host 0.0.0.0: makes the server accessible on all network interfaces
• --port 8080: specifies the port used for incoming requests
• -n 512: sets the default maximum number of generated tokens

Once the server starts successfully, it can be accessed at: http://127.0.0.1:8080.

This allows other applications, scripts, or web interfaces to communicate with Gemma 4 through a local API instead of launching a new inference process for every request.

Enabling Vision Support with MMProj

Although Gemma 4 is a multimodal model, image understanding requires an additional MMProj model when running through llama.cpp.

The MMProj model converts visual features produced by the vision encoder into the embedding space expected by the language model, allowing Gemma 4 to reason jointly over images and text.

To enable vision support, launch the server with the additional --mmproj argument.

./build/bin/llama-server \
 -m /Users/cosmo3769/.lmstudio/models/lmstudio-community/gemma-4-E2B-it-GGUF/gemma-4-E2B-it-Q4_K_M.gguf \
 --mmproj /Users/cosmo3769/.lmstudio/models/lmstudio-community/gemma-4-E2B-it-GGUF/mmproj-gemma-4-E2B-it-BF16.gguf \
 --host 0.0.0.0 \
 --port 8080 \
 -n 512

Once the MMProj model is loaded, the server can process multimodal requests containing both text and images, enabling tasks such as image captioning, visual question answering, and general image understanding while continuing to use the same GGUF language model.

To learn more about llama.cpp, you can refer to this blog from PyImageSearch.

Running Gemma 4 with MLX on Apple Silicon Macs

If you are using an Apple Silicon Mac, MLX is one of the most efficient ways to run Gemma 4 locally. Developed by Apple, MLX is a machine learning framework built specifically for Apple Silicon, enabling models to take full advantage of the unified memory architecture and GPU acceleration available on M-series chips.

Building on top of MLX, MLX-VLM extends the framework with support for vision-language models, making it straightforward to run multimodal models such as Gemma 4 with just a few commands. It also provides an OpenAI-compatible server, allowing locally hosted models to integrate seamlessly with existing applications and tools that already support the OpenAI API.

In this section, we will install MLX-VLM, perform multimodal inference from the command line, launch a local OpenAI-compatible server, and interact with Gemma 4 using both cURL and the official OpenAI Python SDK (Software Development Kit).

Creating a Python Environment

We begin by creating a dedicated Python virtual environment.

python3 -m venv .venv
source .venv/bin/activate

The first command creates an isolated Python environment named .venv, while the second activates it.

Using a virtual environment keeps project dependencies separate from the system-wide Python installation and helps avoid version conflicts between different machine learning projects.

Installing MLX-VLM

Next, install the MLX-VLM package.

pip install mlx-vlm

The mlx-vlm package includes everything needed to download, load, and run multimodal vision-language models on Apple Silicon devices.

Once installed, we can immediately begin running Gemma 4 without manually downloading model files or compiling inference engines.

Running Single-Image Inference

MLX-VLM provides a convenient command-line interface for performing multimodal inference.

mlx_vlm.generate \
 --model mlx-community/gemma-4-e2b-it-4bit \
 --image https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg \
 --prompt "Describe this image in detail"

Let us examine the command:

• --model: specifies the Gemma 4 checkpoint to load
• --image: points to the image that will be analyzed
• --prompt: provides the accompanying text instruction for the model

If the model is not already available locally, MLX automatically downloads it from Hugging Face before performing inference.

After loading the image and prompt, Gemma 4 processes both inputs together and generates a detailed textual description of the image.

Launching an OpenAI-Compatible Server

Instead of running one inference at a time, we can launch a local server that exposes an OpenAI-compatible API.

mlx_vlm.server \
 --model mlx-community/gemma-4-e2b-it-4bit \
 --port 8080

This command loads the model into memory and starts a local inference server listening on port 8080.

Once running, any application capable of communicating with the OpenAI Chat Completions API can send requests directly to the local server without requiring changes to the overall application logic.

Querying the Server with cURL

With the server running, we can send requests using standard HTTP calls.

curl http://localhost:8080/v1/chat/completions \
 -H "Content-Type: application/json" \
 -d '{
   "model": "mlx-community/gemma-4-e2b-it-4bit",
   "messages": [
     {
       "role": "user",
       "content": [
         {"type": "text", "text": "Describe this image in detail"},
         {"type": "image_url", "image_url": {"url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg"}}
       ]
     }
   ],
   "max_tokens": 500
 }'

The request closely mirrors the OpenAI Chat Completions API.

Unlike text-only requests, the content field now contains a list consisting of both:

• a text instruction
• an image URL

Gemma 4 processes the image and prompt together before generating a multimodal response.

The max_tokens parameter specifies the maximum number of tokens the model may generate.

Querying the Server with Python and OpenAI SDK

Since MLX-VLM exposes an OpenAI-compatible endpoint, we can also interact with it using the official OpenAI Python SDK.

First, install the SDK:

pip install openai --break-system-packages

Next, create an OpenAI client that points to the locally running server.

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8080/v1", api_key="fake")

Notice that the API key is simply a placeholder. Authentication is not required because the server is running locally on our machine.

Once the client has been initialized, sending multimodal requests is almost identical to interacting with OpenAI-hosted models.

response = client.chat.completions.create(
   model="mlx-community/gemma-4-e2b-it-4bit",
   messages=[
       {
           "role": "user",
           "content": [
               {"type": "text", "text": "Describe this image"},
               {"type": "image_url", "image_url": {"url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/0052a70beed5bf71b92610a43a52df6d286cd5f3/diffusers/rabbit.jpg
"}}
           ]
       }
   ],
   max_tokens=500
)

The request contains both text and image inputs inside the messages field, following the same structure used by the OpenAI Vision API.

After the request is sent, Gemma 4 analyzes the image, combines it with the accompanying prompt, and returns a textual response, which we print using:

print(response.choices[0].message.content)

One of the biggest advantages of MLX-VLM is that developers already using the OpenAI SDK can switch to a fully local Gemma 4 deployment with minimal code changes. In many cases, updating the base_url to point to the local server is all that is required, making it straightforward to move between cloud-hosted and on-device inference while keeping the rest of the application unchanged.

Running Gemma 4 with LM Studio’s Desktop App

If you are looking for the easiest way to run Gemma 4 locally without using the command line, LM Studio is an excellent choice. It provides an intuitive graphical interface for downloading, managing, and interacting with local language models, while handling much of the underlying configuration automatically.

Behind the scenes, LM Studio uses llama.cpp as its inference engine for GGUF models, allowing you to benefit from efficient CPU and GPU inference without manually compiling the runtime or configuring model files. This makes it an ideal option for developers who want a streamlined setup while still taking advantage of the performance offered by llama.cpp.

In this section, we will install LM Studio, download a Gemma 4 model, configure the local runtime, and begin interacting with Gemma 4 through its built-in chat interface.

Installing LM Studio

Begin by downloading and installing LM Studio from its official website. Once the installation is complete, launch the application.

On the first launch, you will be greeted with the welcome screen shown in Figure 18.

Click Get Started to begin the setup process.

Downloading Your First Model

LM Studio guides you through downloading your first local model.

As shown in Figure 19, it recommends downloading a Gemma 4 model directly from within the application.

Click Download gemma-4-e4b to begin downloading the model. During the download, LM Studio displays the current progress while allowing you to continue using the application.

Once the download finishes, you will see a confirmation screen indicating that the model is ready for use.

Click Continue to proceed to the final setup step.

Enabling Developer Mode

Before entering the main interface, LM Studio presents several optional developer settings.

For this lesson, enable both options:

• Turn on Developer Mode
• Start local LLM service on login

Developer Mode exposes additional features useful for developers, while automatically starting the local server makes it easier to access your models from external applications through an OpenAI-compatible API.

After enabling these options, click Continue to LM Studio.

Loading Gemma 4

Once LM Studio opens, you will see an empty workspace similar to Figure 23.

To begin chatting with Gemma 4, create a new conversation by clicking New Chat.

Next, click Select a model to load at the top of the window.

If you have already downloaded Gemma 4 during setup, you can simply select it from the available models.

Otherwise, use the search bar to search for gemma-4-e2b. LM Studio displays several compatible model variants, including GGUF and MLX versions.

For this lesson, select the gemma-4-E2B-it-GGUF model and click Download. LM Studio automatically downloads the model and stores it in its local model library.

After the download completes, load the model into memory by selecting it from the model picker.

Once loaded successfully, the model name appears at the top of the interface, and the chat window becomes active.

Running Inference

You can now interact with Gemma 4 just like any other chat assistant. Simply enter your prompt into the message box and press Enter to generate a response.

For example, asking:

What is an LLM?

produces the response shown in Figure 28.

In Figure 29, we can also see that we can input an image and prompt together to generate a response.

At this point, we have a fully functional local Gemma 4 environment capable of running entirely on our own machine, without relying on external APIs or cloud services.

To learn more about LM Studio, you can refer to this blog from PyImageSearch.

Running Gemma 4 Locally with Transformers.js

Unlike the deployment options we have explored so far, Transformers.js allows us to run Gemma 4 entirely from JavaScript. Instead of relying on Python or a native inference engine, it enables developers to load and run Hugging Face models directly within Node.js applications using familiar JavaScript APIs.

Under the hood, Transformers.js executes ONNX versions of Hugging Face models through ONNX Runtime, making it possible to perform local inference without leaving the JavaScript ecosystem. This makes it an excellent choice for JavaScript developers who want to build AI-powered applications without introducing a separate Python backend or external inference service.

In this section, we will install Transformers.js, download an ONNX version of Gemma 4, load the model locally, process both image and audio inputs, and perform multimodal inference entirely from JavaScript.

Setting Up the Project

We begin by creating a new Node.js project.

npm init -y

This command initializes a new Node.js project and creates a package.json file containing the project’s metadata, dependencies, and configuration.

npm install @huggingface/transformers wavefile

Next, install the required packages.

This installs 2 packages that our application depends on:

• @huggingface/transformers: provides the Transformers.js library for downloading, loading, and running Hugging Face models directly from JavaScript.
• wavefile: is a lightweight library for reading and preprocessing WAV audio files before they are passed to Gemma 4.

Once the installation completes, create a file named index.mjs. This file will contain our entire multimodal inference pipeline, including model loading, input preprocessing, and response generation.

Importing the Required Libraries

We begin by importing the libraries needed throughout the application.

import {
 AutoProcessor,
 Gemma4ForConditionalGeneration,
 TextStreamer,
 load_image,
} from "@huggingface/transformers";
import pkg from "wavefile";
const { WaveFile } = pkg;
import { readFileSync, existsSync, createWriteStream } from "fs";
import { pipeline } from "stream/promises";
import { tmpdir } from "os";
import { join } from "path";

Let us understand the purpose of each import.

The first group of imports comes from Transformers.js.

• AutoProcessor: loads the preprocessing pipeline associated with Gemma 4. It prepares text, images, and audio before they are passed to the model.
• Gemma4ForConditionalGeneration: loads the Gemma 4 ONNX model and provides the generation API that we will use for multimodal inference.
• TextStreamer: streams generated tokens to the terminal as they are produced instead of waiting for the entire response to finish. This creates a more interactive experience, especially for longer generations.
• load_image: simplifies loading images from either local files or remote URLs.

Next, we import the wavefile package.

Unlike images, audio inputs require additional preprocessing before they can be passed to Gemma 4. The WaveFile class allows us to read WAV (Waveform Audio File Format) files, convert them to the required bit depth, resample them to the appropriate sampling rate, and extract the raw audio samples expected by the model.

Finally, we import several built-in Node.js modules. These utilities help us manage files and temporary resources during inference.

Here:

• fs: provides functions for reading files from disk and creating temporary files
• pipeline: enables efficient streaming of downloaded data, such as audio files, directly into a file
• tmpdir: returns the operating system’s temporary directory, which we use for storing downloaded files during preprocessing
• join: constructs platform-independent file paths, ensuring the script works consistently across different operating systems

Together, these libraries provide everything needed to load Gemma 4, preprocess multimodal inputs, and perform inference entirely from a JavaScript application.

Creating Helper Functions

Before loading the model, let us define a few helper functions that make our script easier to read and improve the overall user experience.

These functions are not directly involved in running Gemma 4. Instead, they provide utilities for formatting terminal output, displaying download progress, measuring execution time, and preprocessing audio before it is passed to the model.

Formatting Console Output

We begin by defining several ANSI (American National Standards Institute) escape codes.

const RESET  = "\x1b[0m";
const BOLD   = "\x1b[1m";
const GREEN  = "\x1b[32m";
const CYAN   = "\x1b[36m";
const YELLOW = "\x1b[33m";
const DIM    = "\x1b[2m";

ANSI escape codes allow us to style text printed to the terminal. Throughout the script, we use them to distinguish different types of messages by applying colors and formatting.

For example:

• Cyan: is used for informational messages.
• Green: indicates successful operations.
• Yellow: highlights warnings.
• Bold: makes important messages easier to identify.
• Dim: is used when displaying elapsed execution times.

Although optional, these formatting codes make the script much easier to follow, especially since downloading and loading large language models can take several minutes.

Logging Progress

Next, we define 3 small helper functions.

function log(label, msg = "") {
 console.log(`${BOLD}${CYAN}[${label}]${RESET} ${msg}`);
}

function success(label, msg = "") {
 console.log(`${BOLD}${GREEN}[${label}]${RESET} ${msg}`);
}

function warn(msg) {
 console.log(`${BOLD}${YELLOW}[WARN]${RESET} ${msg}`);
}

Rather than repeatedly writing lengthy console.log() statements throughout the program, these helper functions provide a consistent way to display messages.

For example:

• log(): prints general progress updates
• success(): indicates that a step completed successfully
• warn(): displays warning messages that the user should be aware of

Using dedicated helper functions also keeps the main inference pipeline much cleaner and easier to read.

Measuring Execution Time

The next helper measures how long each stage of the pipeline takes to execute.

function elapsed(start) {
 return `${DIM}(${((Date.now() - start) / 1000).toFixed(1)}s)${RESET}`;
}

This function accepts a starting timestamp, computes the elapsed time using Date.now(), converts the result to seconds, and formats it for display.

We will use this helper throughout the script to report how long it takes to load the processor, initialize the model, preprocess inputs, and generate responses.

Displaying Model Download Progress

The first time the script runs, Transformers.js must download the ONNX (Open Neural Network Exchange) model files from Hugging Face. Depending on your network connection and the model size, this process may take several minutes.

To provide better feedback during the download, we define a helper function that renders a progress bar inside the terminal.

// Single overall progress bar that updates in place
let progressStarted = false;
function renderOverallProgress(pct) {
 const clamped = Math.min(Math.round(pct ?? 0), 100);
 const filled = Math.round(clamped / 2);
 const bar = "█".repeat(filled) + "░".repeat(50 - filled);
 if (!progressStarted) {
   process.stdout.write("\n");
   progressStarted = true;
 }
 process.stdout.write(`\r  Downloading  [${CYAN}${bar}${RESET}] ${clamped.toString().padStart(3)}%`);
 if (clamped >= 100) process.stdout.write(`  ${GREEN}✓${RESET}\n`);
}

The function receives the current download percentage, converts it into a visual progress bar, and continuously updates the same terminal line as the download progresses.

Instead of printing hundreds of individual progress messages, the user sees a single progress bar that gradually fills until the download completes.

This provides a much cleaner and more user-friendly experience while large model files are being retrieved.

Loading and Preprocessing Audio

Finally, we define the loadAudioNode() helper.

// Load WAV audio manually for Node.js (no AudioContext needed)
async function loadAudioNode(url) {
 const tmpPath = join(tmpdir(), "audio_input.wav");
 const res = await fetch(url);
 if (!res.ok) throw new Error(`Failed to fetch audio: ${res.statusText}`);
 const dest = createWriteStream(tmpPath);
 await pipeline(res.body, dest);
 const wav = new WaveFile(readFileSync(tmpPath));
 wav.toBitDepth("32f");
 wav.toSampleRate(16000);
 let samples = wav.getSamples();
 // If stereo, use only the first channel
 if (Array.isArray(samples)) samples = samples[0];
 return new Float32Array(samples);
}

Unlike images, audio requires several preprocessing steps before it can be passed into Gemma 4.

This helper performs the entire pipeline automatically.

Specifically, it:

• downloads the audio file from the provided URL
• saves it to a temporary location on disk
• opens the WAV file using the WaveFile library
• converts the audio to 32-bit floating-point samples
• resamples the audio to 16 kHz, which is the sampling rate expected by Gemma 4
• If the audio contains multiple channels (e.g., stereo), it keeps only the first channel and converts the result into a Float32Array.

The returned Float32Array contains the raw waveform samples that will later be passed to the processor alongside the text prompt and image.

By moving all of this logic into a dedicated helper function, the main inference pipeline remains concise and focused on the high-level workflow rather than the details of audio preprocessing.

Loading the Gemma 4 Processor

With the helper functions in place, we are ready to load the processor and the Gemma 4 model.

We begin by specifying the model that we will use throughout this lesson.

const model_id = "onnx-community/gemma-4-E2B-it-ONNX";

Unlike the previous lessons that used Hugging Face Transformers in Python, Transformers.js works with ONNX models. Here, we are loading the Gemma 4 E2B Instruct model that has been converted to the ONNX format and published by the Hugging Face community.

Using a single model_id variable also makes it easy to switch to another compatible model later by changing just one line of code.

// 1. Processor
log("1/5", `Loading processor from ${BOLD}${model_id}${RESET} ...`);
let t = Date.now();
const processor = await AutoProcessor.from_pretrained(model_id);
success("1/5", `Processor ready ${elapsed(t)}`);

Next, we load the processor.

The processor is responsible for converting our multimodal inputs into the format expected by Gemma 4.

Depending on the task, it handles operations such as:

• tokenizing text
• preprocessing images
• extracting audio features
• packaging all modalities into tensors that can be passed directly to the model

Calling AutoProcessor.from_pretrained() automatically downloads the processor configuration (if it is not already available locally) and initializes the preprocessing pipeline associated with the selected model.

We also record the start time using Date.now() so that our helper function can report how long the processor takes to load.

Loading the Gemma 4 Model

After the processor has been initialized, we can load the Gemma 4 model itself.

// 2. Model
log("2/5", `Loading model (dtype=q4f16, device=cpu) — this may take several minutes on first run ...`);
warn("The ONNX model files are large. Progress bar will appear below as files download/load.");
t = Date.now();

const model = await Gemma4ForConditionalGeneration.from_pretrained(model_id, {
 dtype: "q4f16",
 device: "cpu",
 progress_callback: (info) => {
   if (info.status === "progress_total") {
     renderOverallProgress(info.progress);
   }
   // suppress all other per-file noise
 },
});
success("2/5", `Model loaded ${elapsed(t)}`);

Unlike the processor, loading the model can take noticeably longer because Transformers.js may need to download several large ONNX weight files before inference can begin.

Let us look at the most important arguments passed to from_pretrained():

• dtype: "q4f16": loads a 4-bit quantized version of the model. Quantization significantly reduces memory usage and improves inference speed while maintaining good generation quality.
• device: "cpu": instructs Transformers.js to execute the model on the CPU. Depending on your environment, other execution backends may also be available.
• progress_callback: receives updates while the model files are being downloaded and loaded. Instead of displaying numerous progress messages, we pass the overall download percentage to our renderOverallProgress() helper, which renders a clean progress bar in the terminal.

Once the model has been downloaded, subsequent executions become much faster because Transformers.js reuses the locally cached files.

Building the Multimodal Prompt

With both the processor and model loaded, the next step is defining the conversation that will be sent to Gemma 4.

// 3. Prompt
log("3/5", "Building chat prompt ...");
const messages = [
 {
   role: "user",
   content: [
     { type: "image" },
     { type: "audio" },
     {
       type: "text",
       text: "Describe this image in detail and transcribe this audio verbatim.",
     },
   ],
 },
];

Just like the Python Transformers library, Transformers.js represents conversations as a list of chat messages.

Here, the user message contains 3 different modalities:

• an image
• an audio clip
• a text instruction

The text prompt asks Gemma 4 to perform 2 tasks simultaneously:

• describe the contents of the image
• transcribe the accompanying audio

Notice that the image and audio are represented using placeholder objects. The actual image and audio data will be loaded and attached in the next step.

Next, we convert these messages into the format expected by Gemma 4.

const prompt = processor.apply_chat_template(messages, {
 enable_thinking: false,
 add_generation_prompt: true,
});
success("3/5", "Prompt ready");

The apply_chat_template() method formats the conversation according to Gemma 4’s chat template.

Here:

• enable_thinking: false: disables the model’s thinking mode so that only the final response is generated.
• add_generation_prompt: true: appends the appropriate assistant generation token, indicating where the model should begin generating its response.

Using the model’s built-in chat template ensures that prompts follow the exact conversational format expected during training.

Loading and Processing the Inputs

Next, we load the image and audio that will be provided to the model.

// 4. Inputs
log("4/5", "Fetching image and audio, then processing inputs ...");
t = Date.now();

const imageUrl = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/artemis.jpeg";
const audioUrl = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/jfk.wav";

For this example, both inputs are hosted remotely and downloaded at runtime.

We first load the image.

log("4/5", `  Image → ${imageUrl}`);
const image = await load_image(imageUrl);
success("4/5", "  Image loaded");

The load_image() helper downloads the image and prepares it for the processor.

Next, we load the audio.

log("4/5", `  Audio → ${audioUrl}`);
const audio = await loadAudioNode(audioUrl);
success("4/5", `  Audio loaded — ${audio.length} samples @ 16kHz`);

As discussed earlier, loadAudioNode() downloads the WAV file, converts it to 32-bit floating-point samples, resamples it to 16 kHz, and returns a Float32Array containing the waveform.

With both inputs available, we can prepare the model inputs.

const inputs = await processor(prompt, image, audio, { add_special_tokens: false });
success("4/5", `Inputs ready ${elapsed(t)}`);

The processor combines the formatted prompt, image, and audio into a single set of tensors suitable for Gemma 4.

Notice that all 3 modalities are processed together through a single API call. This is one of the advantages of Transformers.js, where the same processor handles text, images, and audio without requiring separate preprocessing pipelines.

Running Inference

With the inputs prepared, we can finally generate a response.

// 5. Generate
log("5/5", "Running inference (CPU can be slow — please wait) ...");
warn("Token generation will stream below as it completes:\n");
t = Date.now();

let tokenCount = 0;
const outputs = await model.generate({
 ...inputs,
 max_new_tokens: 512,
 do_sample: false,
 streamer: new TextStreamer(processor.tokenizer, {
   skip_prompt: true,
   skip_special_tokens: false,
   callback_function: (text) => {
     tokenCount++;
     process.stdout.write(text);
   },
 }),
});

process.stdout.write("\n\n");
success("5/5", `Generation done — ${tokenCount} tokens ${elapsed(t)}`);

The generate() method performs autoregressive text generation using the prepared multimodal inputs.

Some of the most important parameters are:

• max_new_tokens: 512: limits the maximum length of the generated response.
• do_sample: false: disables random sampling, producing deterministic outputs for identical inputs.
• streamer: enables token streaming so that text appears in the terminal as soon as it is generated.

Without a streamer, the program would wait until generation finishes before displaying the complete response. By using TextStreamer, users receive immediate feedback as Gemma 4 produces each token, resulting in a much more interactive experience.

Decoding the Generated Output

The output returned by generate() consists of token IDs rather than human-readable text.

To convert those tokens back into natural language, we decode the generated sequence.

// Final decoded output
const decoded = processor.batch_decode(
 outputs.slice(null, [inputs.input_ids.dims.at(-1), null]),
 { skip_special_tokens: true },
);

console.log(`\n${BOLD}── Final Output ──────────────────────────────────${RESET}`);
console.log(decoded[0]);
console.log(`${BOLD}──────────────────────────────────────────────────${RESET}\n`);

Here, we first remove the original prompt tokens, keeping only the newly generated portion of the sequence. We then call batch_decode() to transform the token IDs into readable text while removing any special control tokens.

Finally, we print the generated response to the terminal.

At this point, Gemma 4 has successfully processed the text prompt, image, and audio inputs, producing a unified multimodal response, all from a single JavaScript application running locally through Transformers.js.

Output

In your terminal, run node index.mjs to run inference using Gemma 4 and Transformers.js.

Choosing the Right Runtime for Running Gemma 4 Locally

Throughout this lesson, we have explored five different ways to run Gemma 4 locally. While each runtime supports local inference, they target different workflows and use cases. The right choice depends on your goals, preferred development environment, and hardware.

Choose Ollama if…

Use Ollama when you want the quickest way to get started with Gemma 4.

It is ideal for developers who:

• want a simple installation process
• prefer a command-line interface with minimal configuration
• need an OpenAI-compatible API for existing applications
• want to experiment with local models in just a few commands

For most developers getting started with local LLMs, Ollama is the easiest recommendation.

Choose llama.cpp if…

Use llama.cpp when you need maximum control and performance.

It is a good fit if you:

• want to work directly with GGUF models
• need fine-grained control over inference settings
• plan to optimize CPU or GPU performance
• want to build custom local inference pipelines

Although it requires more setup than Ollama or LM Studio, llama.cpp offers the greatest flexibility.

Choose MLX if…

Use MLX if you are working on an Apple Silicon Mac.

It is particularly useful when you:

• own an M-series Mac
• want native Apple Silicon performance
• need efficient multimodal inference
• prefer using an OpenAI-compatible local server

For Apple users, MLX typically provides the best performance with the least amount of configuration.

Choose LM Studio if…

Use LM Studio when you prefer a graphical user interface over the command line.

It is well suited for users who:

• want to download and manage models visually
• prefer chatting with models through a desktop application
• need a local API server without compiling or configuring runtimes
• are new to local AI and want the simplest desktop experience

LM Studio is an excellent choice for beginners or anyone who prefers a GUI-based workflow.

Choose Transformers.js if…

Use Transformers.js when you are building JavaScript applications.

It is the best option if you:

• develop with Node.js
• want to perform inference without Python
• need to integrate Gemma 4 directly into a JavaScript project
• plan to eventually deploy models inside the browser

Since Transformers.js powers both Node.js and browser-based inference, it provides a natural path from local development to fully client-side AI applications.

At a Glance

Ultimately, there is no single “best” runtime. Each excels in different scenarios. If you are just getting started, Ollama or LM Studio are excellent entry points. If you need maximum flexibility, llama.cpp is hard to beat. Apple Silicon users should strongly consider MLX, while JavaScript developers will likely feel most at home with Transformers.js.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: August 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, we explored 5 ways to run Gemma 4 locally, each designed to support a different development workflow.

We started with Ollama, the quickest way to download models, launch local inference, and expose an OpenAI-compatible API using just a few commands. We then moved to llama.cpp, where we built the runtime from source, downloaded GGUF models, enabled hardware acceleration, and launched both command-line and API-based inference.

Next, we explored MLX, Apple’s machine learning framework for Apple Silicon, and used MLX-VLM to perform multimodal inference while exposing an OpenAI-compatible server. We then looked at LM Studio, which provides a user-friendly graphical interface for downloading, managing, and interacting with Gemma 4 without requiring any command-line configuration.

Finally, we stepped into the JavaScript ecosystem with Transformers.js, demonstrating how to load an ONNX version of Gemma 4, process text, images, and audio, and perform multimodal inference entirely from a Node.js application without relying on Python.

As we have seen throughout this lesson, there is no single “best” way to run Gemma 4. The ideal runtime depends on your workflow, hardware, and application requirements. Ollama offers the fastest path to getting started, llama.cpp delivers maximum flexibility and performance, MLX is an excellent choice for Apple Silicon devices, LM Studio provides an intuitive desktop experience, and Transformers.js enables JavaScript developers to integrate Gemma 4 directly into their applications.

While all of these approaches run Gemma 4 locally, they still execute as traditional desktop or server-side applications. But what if we could eliminate even that requirement and run Gemma 4 entirely inside a web browser?

In the next lesson, we will build on the Transformers.js foundation introduced here and take the next step by running Gemma 4 directly in the browser. We will learn how to load ONNX models with Transformers.js, leverage browser technologies such as WebGPU and WebAssembly for local inference, and build fully client-side multimodal AI applications that require no Python backend or dedicated inference server.

See you in the next lesson!

Citation Information

Thakur, P. “Running Gemma 4 Locally: Ollama, llama.cpp, MLX, and More,” PyImageSearch, S. Huot, G. Kudriavtsev,and A. Sharma, eds., 2026, https://pyimg.co/1rpad

@incollection{Thakur_2026_running-gemma-4-locally-ollama-llama-cpp-mlx,
  author = {Piyush Thakur},
  title = {{Running Gemma 4 Locally: Ollama, llama.cpp, MLX, and More}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Georgii Kudriavtsev and Aditya Sharma},
  year = {2026},
  url = {https://pyimg.co/1rpad},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post Running Gemma 4 Locally: Ollama, llama.cpp, MLX, and More appeared first on PyImageSearch.

Building Multimodal AI Applications with Gemma 4 and Transformers

In the first part of this series, we explored the architecture behind Gemma 4: the interleaved attention design, Mixture-of-Experts routing, multimodal encoders, Per-Layer Embeddings, and the engineering decisions that allow the model family to scale from smartphones to large GPU servers. We also looked at the different variants, hardware requirements, benchmarks, and the reasoning capabilities that make Gemma 4 one of the most compelling open-weight model releases today.

But architecture alone does not tell the full story.

The real question for most developers is simple: how do we actually use Gemma 4 in practice?

That is what this second part focuses on. Instead of discussing theory, we will build directly with the model using Hugging Face Transformers. We will see how Gemma 4 handles images, videos, audio, structured tool calling, object detection, screenshot-to-code generation, and multimodal reasoning, all from a single unified interface.

One of the most interesting things about Gemma 4 is that the workflow stays surprisingly consistent across tasks. Whether we are passing a webpage screenshot, an audio clip, or a video file, the overall pipeline barely changes:

• Define the multimodal message.
• Process it using the Gemma processor.
• Generate outputs from the model.
• Decode and parse the response.

Once this pattern becomes familiar, building multimodal applications starts feeling much more approachable.

In this lesson, we will go section by section through the notebook and explain not just what the code is doing, but why each step matters. Along the way, we will also connect the implementation back to the architectural ideas from Part 1, showing how concepts (e.g., multimodal encoders, thinking mode, and structured outputs) appear in real inference pipelines.

This lesson is the 2nd in a 5-part series on Google DeepMind’s Gemma 4:

1. Google DeepMind’s Gemma 4: MoE, Efficiency Tricks, and Benchmarks
2. Building Multimodal AI Applications with Gemma 4 and Transformers (this tutorial)
3. Lesson 3
4. Lesson 4
5. Lesson 5

To learn how to build multimodal applications with Gemma 4, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Configuring Your Development Environment

To follow this guide, you need to have the following libraries installed on your system.

!pip install -q -U transformers

This command installs the latest version of the transformers library.

Here:

• -U upgrades the package to the newest available version
• -q enables quiet mode to reduce unnecessary installation logs

Using the latest version is important because Gemma 4 support, multimodal processors, and newer generation features are included in recent releases of the library.

The transformers package gives us access to:

• Pretrained Gemma 4 checkpoints
• Multimodal processors
• Tokenization utilities
• Chat templates
• Text generation APIs

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Installing Python Dependencies and Importing Gemma 4 Multimodal Libraries

After configuring the environment, the next step is importing all the required libraries. These imports provide everything needed for loading Gemma 4, processing multimodal inputs, generating outputs, and visualizing results.

import re
import json
import torch
import matplotlib.pyplot as plt
import matplotlib.patches as patches

from PIL import Image
from transformers.image_utils import load_image
from transformers import AutoModelForMultimodalLM, AutoProcessor

We first import Python’s built-in re and json libraries. The re library is used for regular expression operations. Later in the lesson, we use it while parsing structured outputs such as bounding box predictions returned by the model.

The json library helps us work with JSON responses generated by Gemma 4. This becomes especially useful for tasks (e.g., object detection and function calling), where the model produces structured outputs instead of plain text.

Next, we import PyTorch. PyTorch is the deep learning framework powering the model inference pipeline. Hugging Face Transformers internally relies on PyTorch tensors for:

• Token representations
• Model weights
• GPU computation
• Generation operations

Almost every operation inside Gemma 4 inference eventually runs through PyTorch.

We then import Matplotlib utilities. These libraries are used for visualization. In the object detection section later in the lesson, we will draw predicted bounding boxes on images using rectangle patches.

Next, we import PIL. PIL (Python Imaging Library) is widely used for image loading and manipulation in Python applications. It helps us work with image files before sending them into the model.

We also import load_image. This utility simplifies image loading directly from URLs or local paths. Instead of manually downloading and preprocessing images, we can load them with a single function call.

Finally, we import the 2 most important components from Transformers.

AutoModelForMultimodalLM loads the Gemma 4 multimodal model itself. Since Gemma 4 can process text, images, video, and audio, we use the multimodal model interface instead of a text-only causal language model.

AutoProcessor handles preprocessing for all supported modalities. It prepares:

• Text tokens
• Image embeddings
• Audio features
• Video inputs
• Chat templates

into the format expected by Gemma 4.

Together, these imports form the foundation for the rest of the lesson. Once these libraries are loaded, we are ready to initialize the model and begin multimodal inference.

Loading the Gemma 4 Multimodal Model with Hugging Face Transformers

Now that the environment and dependencies are ready, we can load the Gemma 4 model and its processor.

model_id = "google/gemma-4-E2B-it"
model = AutoModelForMultimodalLM.from_pretrained(model_id, device_map="auto")
processor = AutoProcessor.from_pretrained(model_id)

We first define the model checkpoint: google/gemma-4-E2B-it.

Here, we are using the instruction-tuned Gemma-4-E2B-it variant. The it suffix stands for instruction tuned, meaning the model has been optimized for conversational and task-following behavior.

The Gemma 4 E2B model is also one of the smaller Gemma 4 variants, making it practical for experimentation and multimodal inference without requiring extremely large GPU memory. As discussed in the first part of this series, the E-series models are specifically designed for efficient deployment while still supporting text, image, video, and audio understanding.

Next, we load the model itself. The from_pretrained() method downloads the pretrained weights directly from Hugging Face and initializes the model architecture automatically.

The important argument here is: device_map="auto". This tells Transformers to automatically decide where the model should be loaded:

• GPU if CUDA is available
• CPU otherwise

For larger models, this can also distribute layers across multiple GPUs automatically. In our case, it simplifies deployment because we do not need to manually move tensors or model weights between devices.

Finally, we load the processor. The processor is responsible for preparing multimodal inputs before they are passed into Gemma 4.

This includes:

• Tokenizing text
• Processing images
• Handling audio inputs
• Formatting video data
• Applying the correct chat template

One of the biggest advantages of the processor API is consistency. Whether we send text, images, video, or audio into the model, the processor converts everything into the exact format expected by Gemma 4 internally.

At this point, both the model and processor are fully initialized, and we are ready to begin multimodal inference tasks.

Screenshot-to-Code Generation with Gemma 4 Vision-Language AI

One of the most impressive capabilities of Gemma 4 is multimodal code generation. Instead of giving the model only text prompts, we can provide an image of a webpage and ask it to generate the corresponding HTML code.

We begin by defining the multimodal conversation input:

messages = [
   {
       "role": "user",
       "content": [
           {
               "type": "image",
               "image": "https://github.com/PyImageSearch/assets/blob/main/images/landing-page-1.png?raw=true",
           },
           {"type": "text", "text": "Write HTML code for this page."},
       ],
   }
]

Here, the input follows a chat-style structure. Each message contains:

• A role
• A content field

The content itself is a list because Gemma 4 supports multiple modalities within the same conversation turn.

The first content item is the webpage screenshot. This tells the processor that the input modality is an image. Instead of manually downloading and preprocessing the image ourselves, we simply provide the URL.

The second content item is the text instruction. This prompt guides the model toward the task we want it to perform.

Together, the image and text become a unified multimodal input. Gemma 4 processes both inputs simultaneously, allowing it to reason about the webpage layout visually while generating HTML code as output.

inputs = processor.apply_chat_template(
   messages,
   tokenize=True,
   return_dict=True,
   return_tensors="pt",
   add_generation_prompt=True,
   enable_thinking=True,
).to(model.device)

Next, we prepare the inputs using the processor. This is one of the most important steps in the entire workflow.

The apply_chat_template() function converts the multimodal conversation into the exact format expected by Gemma 4 internally.

Let us break down the major arguments.

• tokenize=True: converts the processed input into tokens that the model can understand
• return_tensors="pt": tells the processor to return PyTorch tensors instead of Python lists
• return_dict=True: returns a structured dictionary containing all required tensors instead of returning only token IDs
• add_generation_prompt=True: appends the assistant generation marker so the model knows it should begin generating a response
• enable_thinking=True: activates Gemma 4’s thinking mode. As discussed in Part 1, thinking mode encourages the model to internally reason through the problem before producing the final answer.

Finally, we move all tensors onto the same device as the model. This ensures inference runs correctly on either CPU or GPU.

output = model.generate(**inputs, max_new_tokens=4000)

Once the inputs are ready, we generate the output. The generate() function performs autoregressive decoding and produces the model response token by token. max_new_tokens=4000 sets the maximum generation length. Since HTML generation can produce large outputs, we allow a relatively high token budget.

input_len = inputs.input_ids.shape[-1]
generated_text_ids = output[0][input_len:]
generated_text = processor.decode(generated_text_ids, skip_special_tokens=True)

Next, we extract only the generated portion of the response. This step is necessary because the generated output contains:

• The original input tokens
• The newly generated tokens

We remove the input portion and keep only the generated response. We then decode the generated token IDs back into readable text. The skip_special_tokens=True argument removes internal formatting tokens that are not useful for display.

result = processor.parse_response(generated_text)

print(result["content"])

Next, we parse the model response. This converts the raw generated text into a cleaner structured format. Finally, we print the generated HTML output.

Output

At this stage, Gemma 4 produces HTML code that attempts to recreate the webpage shown in the screenshot.

This example demonstrates one of the most exciting aspects of multimodal language models: the ability to translate visual layouts directly into working code. Instead of treating images and programming as separate domains, Gemma 4 combines visual understanding and code generation into a single inference pipeline.

thought
Here's a thinking process that leads to the suggested HTML code:
1.  **Analyze the Request:** The user wants HTML code for the provided image (a website screenshot/design).
2.  **Deconstruct the Image (Identify Key Components & Structure):**
    *   **Header/Navigation Bar:**
        *   Logo: "YOURLOGO" (Top left).
        *   Navigation Links: "Home," "Company," "Services," "Contact."
        *   Icon: Search/Magnifying glass (Top right).
    *   **Hero Section/Main Content Area:**
        *   *Visual Element (Left/Center):* A stylized illustration of a woman celebrating, surrounded by stars, planets, and dots.
        *   *Text/Headline (Right):*
            *   Headline: "Creative Ideas" (Large, orange/bold).
            *   Subheadline/Tagline: "A Studio based in New York, US."
            *   Body Text: "Solutions of design in modern era, consectetur adipiscing elit, sed diam nonummy nibh euismod tincidunt ut laoreet dolore magna aliquam erat volutpat." (Placeholder text).
        *   *Call to Action (CTA):* A button labeled "SUBSCRIBE" (Dark/Maroon).
        *   *Social Media Icons:* Facebook, Instagram, Twitter (Small icons below the CTA).
...
...

We can see that Gemma 4 not only understands the webpage visually, but also reasons through its structure before generating the final code. The model identifies components such as the navigation bar, hero section, buttons, and social icons, and then converts them into structured HTML and CSS.

Interestingly, the output also includes responsive styling, layout organization, practical implementation notes, and even a dedicated section explaining how to use the generated code. We followed the same instructions provided by Gemma 4 to run the generated HTML and CSS files, and the resulting webpage closely matched the original design shown in the input image.

Video Understanding and Multimodal Reasoning with Gemma 4

After testing image understanding and code generation, we can move one step further and evaluate Gemma 4 on video reasoning tasks.

messages = [
   {
       "role": "user",
       "content": [
           {"type": "video", "url": "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/concert.mp4"},
           {"type": "text", "text": "What is happening in the video? What is the song about?"},
       ],
   },
]

Here, instead of providing an image, we provide a video file as input.

The first content block defines the video source. This tells the processor that the input modality is a video rather than an image or audio clip.

The second content block contains the text instruction. This prompt asks the model to reason about:

• The visual scene
• The event taking place
• The semantic meaning of the song

This is important because the model is not simply captioning individual frames. It is performing multimodal reasoning across both temporal and audio information.

inputs = processor.apply_chat_template(
   messages,
   tokenize=True,
   return_dict=True,
   return_tensors="pt",
   add_generation_prompt=True,
   load_audio_from_video=True,
).to(model.device)

Next, we process the multimodal input. Most of the arguments are similar to the previous example, but there is one important addition here:

• load_audio_from_video=True flag tells the processor to extract the audio stream from the video alongside the visual frames.

As discussed in the first part of this series, the smaller Gemma 4 E-series models support audio understanding in addition to image and text reasoning.

This means the model can jointly analyze:

• Video frames
• Speech
• Music
• Ambient sounds

instead of relying only on visual information.

In practice, this allows Gemma 4 to answer richer questions about videos, such as:

• Understanding spoken dialogue
• Identifying music genres
• Inferring emotional tone
• Explaining events happening in the scene

output = model.generate(**inputs, max_new_tokens=200)

Once the inputs are prepared, we generate the output. Here, max_new_tokens=200 limits the response length. Since video descriptions are usually shorter than HTML generation tasks, a smaller token budget is sufficient.

input_len = inputs.input_ids.shape[-1]
generated_text_ids = output[0][input_len:]
generated_text = processor.decode(generated_text_ids, skip_special_tokens=True)
result = processor.parse_response(generated_text)

After generation, we extract only the generated portion of the sequence. Just as before, this removes the original input tokens and keeps only the newly generated response.

Next, we decode the generated tokens back into readable text. We then parse the response into a cleaner structure.

print(result["content"])

Finally, we print the generated output.

Output

At this stage, Gemma 4 analyzes the concert video and produces a multimodal interpretation of the scene.

This example highlights an important capability of Gemma 4: unified video understanding. Instead of requiring separate models for:

• Video captioning
• Audio transcription
• Scene understanding
• Semantic reasoning

Gemma 4 performs these tasks inside a single multimodal inference pipeline.

Based on the video, here is what is happening:
**What is happening in the video?**
The video captures a live concert performance on a large stage. We see a band performing, with musicians playing instruments (guitars, drums, etc.) under bright stage lights and dramatic blue and white lighting, including significant smoke/fog effects. In the foreground, there is a crowd of people watching the show, with several audience members visible from behind. The energy of the performance seems high, as suggested by the lighting and the engagement of the audience.
**What is the song about?**
The provided lyrics are:
> "I'm falling on the street
> From neck to chest
> Could it be that moments
> Another one I want to be is
> In the storm alone
> I've been all alone
> I never want to see
> Oh this nice sad place you was given fate
> I don't want it

Multimodal Function Calling with Gemma 4

One of the most powerful capabilities introduced in modern language models is function calling. Instead of generating only plain text responses, the model can decide when to invoke external tools and generate structured arguments for them.

In this example, we combine:

• Vision understanding
• Reasoning
• Tool usage

inside a single multimodal workflow.

WEATHER_TOOL = {
   "type": "function",
   "function": {
       "name": "get_weather",
       "description": "Gets the current weather for a specific location.",
       "parameters": {
           "type": "object",
           "properties": {
               "city": {"type": "string", "description": "The city name"},
           },
           "required": ["city"],
       },
   },
}
tools = [WEATHER_TOOL]

We first define a weather tool schema. This structure describes the external tool available to the model. The tool definition contains several important components.

First, we specify the tool type: "type": "function" tells Gemma 4 that the tool represents a callable function.

Next, we define the actual function metadata. "name": "get_weather" is the function name the model will reference when generating tool calls. We also provide a natural language description. This helps the model understand what the tool does and when it should be used.

Next comes the parameter schema. This follows a JSON-style schema format. Here, the function expects a single required argument:

• city

The model will later generate this parameter automatically based on the image and user request. We then place the tool into a list. This allows multiple tools to be passed into the model if needed.

messages = [
   {"role": "user", "content": [
       {"type": "image", "image": "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/venice.jpg"},
       {"type": "text", "text": "What is the city in this image? Check the weather there right now."},
   ]},
]

Next, we define the multimodal user input.

This input combines:

• An image
• A reasoning instruction

The image contains a photograph of Venice, while the text prompt asks the model to:

• Identify the city
• Use the weather tool for that location

This is important because the model must first perform visual reasoning before it can invoke the function correctly.

inputs = processor.apply_chat_template(
   messages,
   tools=[WEATHER_TOOL],
   tokenize=True,
   return_dict=True,
   return_tensors="pt",
   add_generation_prompt=True,
   enable_thinking=True,
).to(model.device)

Next, we prepare the inputs.

The major addition here is: tools=[WEATHER_TOOL]

This injects the available tool schema directly into the prompt context. Gemma 4 now knows:

• Which tools exist
• What each tool does
• Which parameters are required

We also enable reasoning mode again: enable_thinking=True

This helps the model reason through the sequence:

• Analyze image
• Infer city
• Decide tool usage
• Generate function arguments

output = model.generate(**inputs, max_new_tokens=1000)

After processing the inputs, we generate the output. Since function-calling outputs can include reasoning traces and structured responses, we allow a larger token budget.

input_len = inputs.input_ids.shape[-1]
generated_text_ids = output[0][input_len:]
generated_text = processor.decode(generated_text_ids, skip_special_tokens=True)
result = processor.parse_response(generated_text)

Next, we extract only the generated tokens. We then decode the generated response. After decoding, we parse the structured response.

print(result["content"])

Finally, we print the generated output.

Output

At this stage, Gemma 4 identifies the city from the image and generates a structured tool call for retrieving the weather.

This is the input image:

This is the generated output:

thought
Here's a thinking process to arrive at the desired response:
1.  **Analyze the Request:** The user has provided an image and asked two questions:
    *   "What is the city in this image?"
    *   "Check the weather there right now."
2.  **Analyze the Image:**
    *   The image shows a large, ornate building situated on the water, with boats and other buildings in the background.
    *   The architecture strongly resembles famous Italian landmarks, specifically those found in Venice (bridges, canals, specific facade styles).
    *   The presence of a canal setting is a huge clue.
3.  **Identify the Location (Inference/Knowledge Retrieval):**
    *   The architecture points strongly toward Venice, Italy.
    *   The specific building looks like a prominent structure in Venice (e.g., a church or palace facade).
4.  **Determine the Necessary Tool Call:**
    *   The user wants the weather for the identified city.
    *   The available tool is `get_weather(city: STRING)`.
5.  **Formulate the Plan:**
    *   **Step 1:** State the identified city.
    *   **Step 2:** Call the `get_weather` tool with the identified city.

6.  **Execute Step 1 (Internal Check):** The city is Venice.
7.  **Execute Step 2 (Tool Call Generation):**
    *   Call: `get_weather(city="Venice")`
8.  **Final Output Generation:** Since I need to call a tool, I must output the tool call first. (Self-Correction: I cannot give the weather *yet*; I must request the tool execution.)call:get_weather{city:Venice}

This example demonstrates why multimodal function calling is so powerful. The model is no longer limited to passive question answering. Instead, it can:

• Observe visual information
• Reason about the task
• Select external tools
• Generate structured arguments

This forms the foundation for agentic AI systems where multimodal models interact with application programming interfaces (APIs), databases, and external applications dynamically.

Object Detection and Visual Grounding with Gemma 4

Beyond captioning and reasoning, Gemma 4 can also perform structured visual localization tasks such as object detection and pointing. Instead of generating only natural language descriptions, the model can return bounding box coordinates for objects inside an image.

image_url = "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/bird.png"
image = load_image(image_url)

We begin by loading the image. Here, load_image() downloads the image directly from the URL and converts it into a format suitable for further processing.

def resize_to_48_multiple(image):
   w, h = image.size
   new_w = (w // 48) * 48
   new_h = (h // 48) * 48
   return image.crop((0, 0, new_w, new_h))

Next, we define a preprocessing helper function. This function adjusts the image dimensions so both height and width become multiples of 48.

We first extract the original image dimensions. Next, we compute the nearest lower multiples of 48. The // operator performs integer division. This effectively rounds the dimensions downward to the nearest valid multiple.

Finally, we crop the image. This preprocessing step helps ensure compatibility with the vision encoder and patch processing pipeline used internally by Gemma 4. Since transformer-based vision models often process images in fixed patch sizes, maintaining aligned dimensions simplifies inference.

def inputs_for_object_detection(image, what_object):
   messages = [
       {
           "role": "user", "content": [
               {"type": "image", "image": image},
               {"type": "text", "text": f"What's the bounding box for the {what_object} in the image?"}
           ]
       }
   ]

   inputs = processor.apply_chat_template(
       messages,
       tokenize=True,
       add_generation_prompt=True,
       return_dict=True,
       return_tensors="pt",
       enable_thinking=False,
   )

   return inputs.to(model.device)

Next, we define the function that prepares the object detection prompt. This function takes:

• The input image
• The target object name

and constructs the multimodal request.

Inside the function, we first define the chat-style message. The image is passed directly as part of the multimodal input.

The text prompt dynamically inserts the target object name using an f-string. For example, if:

what_object = "bird"

the final prompt becomes:

"What's the bounding box for the bird in the image?"

This allows the same pipeline to work for many different object categories without changing the core logic.

Next, we process the multimodal request. Most of the arguments are familiar from earlier examples, but one important detail here is: enable_thinking=False

Unlike reasoning-heavy tasks (e.g., code generation or function calling), object detection is primarily a localization task. Disabling thinking mode helps keep the output concise and focused on structured bounding box predictions instead of extended reasoning traces.

The processor then:

• Tokenizes the text prompt
• Encodes the image
• Applies the correct chat template
• Converts everything into PyTorch tensors

Finally, we move the tensors onto the same device as the model.

def extract_json(text: str):
   text = text.strip()

   text = re.sub(r"^```(?:json)?\s*", "", text)
   text = re.sub(r"\s*```$", "", text)

   # Try direct parse first
   try:
       return json.loads(text)
   except json.JSONDecodeError:
       pass

   # Fallback: extract first JSON object or array
   match = re.search(r'(\{.*\}|\[.*\])', text, re.DOTALL)
   if match:
       candidate = match.group(1)
       return json.loads(candidate)

   raise ValueError("No valid JSON found")

After preparing the object detection inputs, the next step is running inference, extracting the structured bounding box predictions, and visualizing the detected object on the image.

We first define a helper function for parsing the model response.

The generated output from Gemma 4 is usually returned as text, even when it contains structured JSON. This function cleans and extracts the JSON portion safely.

We first remove unnecessary whitespace. Next, we remove markdown-style code fences if they exist.

Sometimes models return outputs such as:

[
 {
   "box_2d": [...]
 }
]

These regular expressions remove the surrounding Markdown formatting so the remaining content becomes valid JSON.

Next, we attempt direct JSON parsing. If the response is already valid JSON, this step succeeds immediately.

However, model outputs are not always perfectly formatted. Sometimes, additional explanation text appears before or after the JSON block. To handle that, we add a fallback mechanism. This searches for the first JSON object {} or JSON array [] inside the generated response.

If a match is found: we extract and parse only the JSON portion.

Finally, if no valid JSON is detected, the function raises an error.

This helper function is important because structured outputs from multimodal models can sometimes contain extra formatting or reasoning traces. Robust parsing makes downstream processing much more reliable.

def detect_object(image_url, what_object):
   image = load_image(image_url)
   image = resize_to_48_multiple(image)
   inputs = inputs_for_object_detection(image, what_object)
   input_len = inputs["input_ids"].shape[-1]
   generated_outputs = model.generate(**inputs, max_new_tokens=1000, do_sample=False)
   generated = processor.decode(generated_outputs[0, input_len:])
   parsed_json = extract_json(generated)[0]
   return parsed_json

Next, we define the main object detection function. This function performs the complete detection pipeline.

We first load and preprocess the image. Next, we prepare the detection inputs. We then compute the input sequence length. This helps us later separate the generated response from the original prompt tokens.

Next comes inference. Here:

• max_new_tokens=1000: sets the generation limit
• do_sample=False: enables deterministic decoding

Deterministic decoding is useful here because object detection requires stable structured outputs rather than creative variations.

Next, we decode only the generated portion. We then extract the structured JSON response.

The model typically returns a list of detections, so [0] selects the first detection result.

Finally, we return the structured detection dictionary.

def draw_pascal_voc_boxes(i, image, box, label, resize_shape=(1000,1000)):
   dpi = 72
   width, height = image.size
   fig, ax = plt.subplots(1, figsize=[width/dpi, height/dpi], tight_layout={'pad':0})

   ax.imshow(image)

   ymin, xmin, ymax, xmax = box
   re_h, re_w = resize_shape if resize_shape is not None else (height, width)
   xmin = (xmin / re_w) * width
   ymin = (ymin/ re_h) * height
   xmax = (xmax / re_w) * width
   ymax = (ymax/ re_h) * height

   w = xmax - xmin
   h = ymax - ymin

   rect = patches.Rectangle(
       (xmin, ymin),
       w,
       h,
       linewidth=10,
       edgecolor="green",
       facecolor="none"
   )
   ax.add_patch(rect)

   if label is not None:
       ax.text(xmin, ymin-25, label, fontsize=24, bbox=dict(facecolor="yellow", alpha=0.5))

   plt.axis("off")
   plt.savefig(f"boxes_{i}.png")
   plt.close(fig)
   display(fig)

Next, we define the visualization function. This function draws the predicted bounding box on top of the image.

We first create a Matplotlib figure. Then we display the image. Next, we unpack the bounding box coordinates.

Gemma 4 returns normalized coordinates relative to a 1000 × 1000 coordinate grid. We therefore rescale them back to the original image dimensions.

We then compute the bounding box width and height. Next, we create the rectangle overlay. This draws a green bounding box around the detected object. The rectangle is then added onto the image. If a label exists, we also display the object name.

Finally, we save and display the result.

def display_detected_object(image_url, what_object):
   image = load_image(image_url)
   image = resize_to_48_multiple(image)
   detection = detect_object(image_url, what_object)
   box = detection["box_2d"]
   label = detection.get("label", f"{what_object}")
   draw_pascal_voc_boxes("1000", image, box, label)

Next, we define a helper wrapper function. This combines:

• Image loading
• Detection
• Visualization

into one simple pipeline.

We first run object detection. Then extract the returned bounding box and label.

Finally, we visualize the prediction.

display_detected_object("https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/bird.png", "bird")

We can now run the complete pipeline.

Output

At this stage, Gemma 4 identifies the bird in the image, predicts the bounding box coordinates, and overlays the detection visually.

This is the input image:

This is the output image:

This example demonstrates that Gemma 4 is not limited to text generation or captioning. It can also produce structured spatial predictions, enabling workflows such as:

• Object localization
• Visual grounding
• UI element detection
• Document understanding
• Interactive visual reasoning

Image Captioning with Gemma 4 Vision-Language Models

Image captioning is one of the most fundamental multimodal tasks for vision-language models. In this example, we ask Gemma 4 to observe an image and generate a detailed natural language description of the scene.

messages = [
   {
       "role": "user",
       "content": [
           {"type": "image", "url": "https://huggingface.co/datasets/merve/vlm_test_images/resolve/main/mosque.jpg"},
           {"type": "text", "text": "Write single detailed caption for this image."},
       ],
   },
]

We first define the multimodal input message. The input contains 2 parts:

• An image
• A text instruction

The image block specifies the image URL. This tells the processor to load the image and prepare it for the vision encoder inside Gemma 4.

Next, we provide the captioning instruction. The wording of the prompt matters here. Since we request a single detailed caption, the model focuses on generating one coherent descriptive sentence instead of multiple fragmented observations.

inputs = processor.apply_chat_template(
   messages,
   tokenize=True,
   return_dict=True,
   return_tensors="pt",
   add_generation_prompt=True,
).to(model.device)

Next, we process the multimodal input.

This step converts the image and text into the structured format expected by Gemma 4.

The processor internally handles:

• Image preprocessing
• Tokenization
• Prompt formatting
• Tensor creation

output = model.generate(**inputs, max_new_tokens=512)

After preprocessing, we generate the caption. Here, max_new_tokens=512 sets the maximum response length. Captioning tasks generally require fewer tokens than HTML generation or reasoning-heavy workflows, but we still allow enough space for detailed descriptions.

input_len = inputs.input_ids.shape[-1]
generated_text_ids = output[0][input_len:]
generated_text = processor.decode(generated_text_ids, skip_special_tokens=True)
result = processor.parse_response(generated_text)

Next, we isolate the generated response. This removes the original prompt tokens and keeps only the newly generated caption. We then decode the generated tokens into readable text. The response is then parsed into a cleaner structure.

print(result["content"])

Finally, we print the generated caption.

Output

At this stage, Gemma 4 analyzes the image and generates a detailed description of the mosque scene.

This is the input image:

This is the generated output:

A bustling outdoor scene features a grand, domed building, likely a historical or religious structure, surrounded by lush greenery, palm trees, and pedestrians walking along a paved walkway. The sky is bright blue with some white clouds, suggesting pleasant weather. In the background, taller minarets are visible, adding to the architectural grandeur of the location. The foreground shows a metal railing and a crowd of people, indicating a popular public space.

Although captioning appears simpler than function calling or video understanding, it is still a strong test of multimodal reasoning. The model must:

• Identify objects
• Understand spatial relationships
• Infer scene context
• Convert visual information into fluent language

This example also highlights the consistency of the Gemma 4 workflow. Whether we perform:

• Screenshot-to-code generation
• Video reasoning
• Function calling
• Captioning

the overall inference pipeline remains nearly identical. Only the input modality and prompt change.

Audio Understanding with Gemma 4

One of the most interesting capabilities of the Gemma 4 E-series models is native audio understanding. Instead of relying on a separate speech recognition model, Gemma 4 can directly process audio inputs and reason about them inside the same multimodal pipeline.

messages = [
   {
       "role": "user",
       "content": [
           {"type": "audio", "url": "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/winning_call.mp3"},
           {"type": "text", "text": "Can you describe this audio in detail?"},
       ],
   },
]

We begin by defining the multimodal input.

The first content block provides the audio input. This tells the processor that the modality is audio rather than image or video.

The second content block contains the text instruction. This prompt asks the model to analyze and explain the audio content in natural language.

inputs = processor.apply_chat_template(
   messages,
   tokenize=True,
   return_dict=True,
   return_tensors="pt",
   add_generation_prompt=True,
).to(model.device)

Next, we process the multimodal input. This step converts the audio and text into tensors that Gemma 4 can process internally.

As discussed in the first part of this series, the E2B and E4B variants include a dedicated audio encoder. The processor handles the conversion of raw audio into the representation expected by this encoder automatically.

output = model.generate(
   **inputs,
   max_new_tokens=1000,
   do_sample=False,
)

After preprocessing, we generate the response.

Here, max_new_tokens=1000 defines the maximum output length for the generated response.

We also use do_sample=False. This disables sampling and makes generation deterministic. Instead of producing varied outputs each time, the model selects the most likely token at every step. For descriptive tasks (e.g., audio understanding), deterministic decoding often produces more stable and reproducible results.

print(processor.decode(output[0], skip_special_tokens=True))

Finally, we decode and print the generated response.

The decode() function converts token IDs back into readable text, while skip_special_tokens=True removes internal formatting tokens.

Output

At this stage, Gemma 4 analyzes the audio clip and generates a detailed natural language description of what it hears.

This is the generated output:

user
Can you describe this audio in detail?
model
Okay, here's a detailed description of the audio you provided:
**Overall Impression:**
The audio is a lively, energetic, and enthusiastic sports commentary, likely from a baseball game. The tone is excited, building anticipation and celebrating a significant moment.
**Specific Details:**
* **Soundscape:** The audio features the sounds of a live sports broadcast. This includes the voice of a commentator, likely with some background noise of a stadium or crowd, though the focus is clearly on the commentary.
* **Commentary Style:** The commentary is fast-paced, dynamic, and uses typical sports jargon and exclamations. The delivery is high-energy and passionate, reflecting the excitement of the game.
* **Key Phrases and Content:**
    * **"And the O1 pitcher on the way to Edgar Martinez swung on the line."**: This sets the scene, indicating a specific play involving a pitcher and a batter.
    * **"Here comes Joey! Here's Ginger in third base!"**: These are calls to action, identifying players and their positions.
    * **"They've got a way to win!"**: This conveys a sense of hope and determination.
    * **"The problem of plate will be late the manner is glaring the fans for the American League championship!"**: This is a more complex sentence, likely referring to a strategic situation and the importance of the championship.
    * **"I don't believe it!"**: An expression of surprise or disbelief, adding to the excitement.
    * **"And just continues my old time!"**: This suggests a continuation of a successful or exciting sequence of events.
* **Emotional Tone:** The tone is overwhelmingly positive, excited, and celebratory. There's a palpable sense of anticipation and triumph.
**In Summary:**
The audio is a segment of a high-energy baseball game broadcast. It captures a moment of intense action, featuring dynamic commentary that builds excitement and highlights the stakes of a championship game. The commentator's enthusiasm is infectious and effectively conveys the thrill of the moment.

This example highlights an important direction in multimodal AI systems. Traditionally, tasks such as the following:

• Speech recognition
• Audio captioning
• Sound event detection
• Spoken question answering

often required separate specialized models.

Gemma 4 instead handles these tasks within a unified multimodal framework, allowing text, images, video, and audio to all flow through the same inference pipeline.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: August 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, we explored how to build multimodal applications using Gemma 4 and Hugging Face Transformers. Starting from environment setup and model loading, we walked through several real-world multimodal workflows (e.g., including screenshot-to-code generation, video understanding, multimodal function calling, image captioning, audio understanding, and object detection).

One of the biggest takeaways from Gemma 4 is the consistency of its inference pipeline. Whether the input is text, images, video, or audio, the overall workflow remains nearly identical:

• Define the multimodal conversation
• Process the inputs using the processor
• Run generation
• Decode and parse the outputs

This unified design makes multimodal development significantly simpler.

We also saw that Gemma 4 is not limited to conversational AI. It can generate HTML from screenshots, reason over videos, invoke external tools, localize objects with bounding boxes, and understand audio, all inside a single multimodal framework. Together, these capabilities make Gemma 4 a powerful foundation for building next-generation agentic and multimodal AI systems.

Citation Information

Thakur, P. “Building Multimodal AI Applications with Gemma 4 and Transformers,” PyImageSearch, S. Huot, G. Kudriavtsev, and A. Sharma, eds., 2026, https://pyimg.co/09dks

@incollection{Thakur_2026_build-multimodal-ai-apps-w-gemma-4-transformers,
  author = {Piyush Thakur},
  title = {{Building Multimodal AI Applications with Gemma 4 and Transformers}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Georgii Kudriavtsev and Aditya Sharma},
  year = {2026},
  url = {https://pyimg.co/09dks},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post Building Multimodal AI Applications with Gemma 4 and Transformers appeared first on PyImageSearch.

RAG Observability with Langfuse, vLLM, and FAISS

In this lesson, you will learn how to instrument every step of your Retrieval-Augmented Generation (RAG) pipeline using Langfuse, capture traces across ingestion, retrieval, and generation, and understand exactly how your system behaves under the hood.

You will wire tracing into your retriever and generator, monitor latency and token usage, evaluate quality scores, and run the entire stack with vLLM and FAISS locally so you can experiment freely without any cloud dependencies.

By the end, you will have a fully transparent RAG workflow that you can debug, optimize, and scale with confidence.

This lesson is the last in a 3-part series on LLM observability with Langfuse:

1. LLM Observability with Self-Hosted Langfuse and vLLM
2. Manual Tracing, Scores, and Evaluation with Langfuse (Self-Hosted)
3. RAG Observability with Langfuse, vLLM, and FAISS (this tutorial)

To learn how to make your RAG pipeline fully observable with Langfuse, vLLM, and FAISS, just keep reading.

Looking for the source code to this post?

Jump Right To The Downloads Section

Introduction to Production-Grade RAG and LLM Observability

What Makes a RAG Pipeline Production-Grade

A RAG pipeline becomes “production-grade’’ only when it consistently delivers correct, stable, and explainable outputs under real-world constraints. In development, it is easy to get an LLM to answer questions using retrieved context. In production, the challenges multiply: retrieval quality varies, embeddings may shift over time, documents evolve, and latency budgets tighten. A production RAG pipeline must remain robust even when the input data is noisy, queries are unpredictable, and traffic is high.

A production-ready RAG system must treat retrieval as a first-class subsystem, not a background detail. That means surfacing similarity scores, exposing ranking decisions, understanding how vector search behaves at scale, and ensuring retriever recall stays high across diverse query types. It also requires that the prompt construction step is deterministic, inspectable, and traceable, because subtle variations in formatting often change model behavior dramatically.

Beyond these retrieval and prompt concerns, the LLM is also a production component. That means retry logic, token accounting, consistent latency, predictable throughput, and graceful failure modes. Production pipelines need clear boundaries between retrieval failures, prompt-generation bugs, and LLM invocation issues. If these concerns remain invisible, debugging becomes guesswork and reliability collapses under load. Production-grade RAG means engineered behavior, not accidental correctness.

Why Observability Is Essential for Retrieval-Augmented Systems

RAG pipelines fail silently. Retrieval may return irrelevant documents, prompting may omit essential context, and the LLM may hallucinate confidently even when grounded context exists. Without observability, it is impossible to diagnose why a particular answer was wrong. Was the embedding model inconsistent? Did FAISS, the vector search library used to retrieve similar documents, return poor matches? Did the prompt formatting break a system instruction? Did the LLM drift or degrade?

Observability solves this by turning the RAG pipeline into a transparent execution graph. Tools like Langfuse give you hierarchical traces: one trace for the whole request, and nested spans for retrieval, LLM calls, evaluation, and supporting steps. Each span captures inputs, outputs, metadata, latencies, token usage, and even scoring metrics. With this information, problems become diagnosable:

• Retrieval returned low-relevance documents
• Prompt formatting changed unexpectedly
• LLM call degraded or hit retry logic
• Evaluation metrics began trending downward

In other words, observability provides ground truth for system behavior. Production RAG must be accountable: decisions should be explainable, errors traceable, and failures measurable. Without observability, shipping RAG to production is equivalent to flying an airplane without instruments; the system might work, but you will not know when or why it stops working.

What We Will Build: Traced Retriever, Traced LLM, Full RAG Pipeline, and Evaluation

In this lesson, you will construct a fully observable, component-wise traced RAG system using Langfuse, FAISS, SentenceTransformers, and vLLM. Each part of the pipeline is instrumented for visibility: you will build a traced retriever that logs embeddings, index sizes, similarity scores, and ranking. You will build a traced LLM wrapper that records prompts, responses, retry attempts, and token usage. These components power a fully traced RAG pipeline that captures retrieval, prompt construction, generation, and final evaluation as a single hierarchical execution tree.

You will also implement automatic RAG output evaluation, computing relevancy, hallucination risk, and an overall quality score, with each metric logged back to Langfuse dashboards with scoring nodes. This gives you a complete introspection loop: every answer is measured, every metric is recorded, and every decision is traceable through structured spans.

By the end, you will have a production-grade RAG observability stack, running locally with:

• A traced retriever
• A traced LLM client
• A fully instrumented RAG pipeline
• Automatic scoring and diagnostics
• Local dashboards for analyzing behavior

This foundation prepares you for upcoming lessons, where we extend these ideas into multi-step agents and more complex reasoning workflows.

RAG Observability Architecture with Langfuse, vLLM, and FAISS

A production-grade RAG pipeline is not a single model call. It is an orchestrated system composed of independent but cooperating components: retrieval, prompt assembly, LLM inference, and evaluation. In this section, we break down each subsystem and explain how they interact, why they are separated, and how Langfuse stitches everything together into a fully observable execution graph.

Retrieval → Prompt Construction → LLM → Scoring (The Core RAG Loop)

A well-designed RAG architecture follows a clean, linear flow where each stage has a single responsibility:

Retrieval

The system begins by embedding the user query and searching for relevant documents in a vector index. The retriever returns ranked, scored context items that will guide the LLM. In production, retrieval quality is often the primary bottleneck; if retrieval fails, generation cannot be correct. Therefore, retrieval spans log:

• embeddings used
• search distances and converted relevance scores
• number of documents returned
• FAISS query latencies

Your TracedRetriever does exactly this in the codebase, generating embeddings, searching the FAISS index, and tracing each step.

Prompt Construction

Next, the system converts retrieved documents into a structured context block and assembles a prompt that the LLM can reliably parse. Prompt construction must be deterministic to avoid instability across runs. The code in rag_pipeline.py builds a system message, a user message, and contextual references ([1] doc1, [2] doc2, etc.). This ensures:

• deterministic ordering
• visible context structure
• stable interface for downstream evaluation

LLM Generation

The prompt is sent to the LLM via an OpenAI-compatible Completion API, served by vLLM locally. The TracedLLMClient wraps this call with:

• retry logic
• token usage reporting
• error logging
• prompt and response capture
• metadata annotations

This is critical for production reliability because LLM latency, token usage, and intermittent failures must all be observable.

Scoring and Evaluation

Finally, the answer is passed through a lightweight evaluation module (evaluation.py). It computes:

• a relevancy score
• a hallucination risk score
• an overall quality score

These metrics are reported back into Langfuse as scoring nodes. Production RAG systems need this because correctness is subjective; evaluation makes correctness measurable.

This 4-step pipeline forms the backbone of every modern retrieval-augmented system.

Local Vector Store Using FAISS and SentenceTransformers

RAG pipelines must remain fast, private, and cost-efficient. This system uses FAISS as the vector index and SentenceTransformers for embedding models, giving you:

• Zero API cost (everything is local)
• GPU acceleration optional (FAISS works on CPU just fine)
• Deterministic embeddings (critical for reproducibility)
• Config-driven control over the embedding model and dimensionality

The retrieval pipeline is built around the following 3 core mechanisms:

Document Embeddings

Each document is encoded using the local model defined in config.yaml:

embeddings:
  model: "sentence-transformers/all-MiniLM-L6-v2"
  dimension: 384

The TracedRetriever loads this model and produces normalized embeddings for better retrieval precision.

FAISS Index

FAISS stores all document embeddings in a vector index created via:

self.index = faiss.IndexFlatL2(self.dimension)

IndexFlatL2 is simple, fast, and perfect for local development, while still appropriate for many production environments.

Similarity Search

Retrieval happens by computing L2 distance and converting those distances into relevance scores, ensuring consistency and interpretability.

You end up with a fully local, high-performance vector store without touching external cloud APIs.

vLLM as an OpenAI-Compatible Inference Server

Instead of relying on OpenAI or Anthropic APIs, your lesson uses vLLM, a high-throughput inference engine built for serving LLMs at scale.

Your Docker Compose file runs vLLM either on GPU (recommended) or CPU fallback, exposing it at:

http://localhost:8000/v1

This allows you to call vLLM with the exact same interface as OpenAI:

response = client.chat.completions.create(
    model=self.model,
    messages=messages,
    temperature=temperature,
    max_tokens=max_tokens
)

Benefits for production-grade RAG:

• Predictable latency
• Full control over model versioning
• No external dependencies
• High-throughput serving (paged attention)
• OpenAI API compatibility (no code rewrite needed)

Your TracedLLMClient wraps all of this with Langfuse observability, giving you:

• latency metrics
• retry attempts
• token usage breakdown
• full input/output transparency
• error-level spans when inference fails

This is how modern enterprises run private LLMs with production reliability.

Langfuse for Tracing, Metrics, Evaluation, and Span Hierarchies

Langfuse is the backbone of observability in this system. Every major component (i.e., embedding, retrieval, generation, and evaluation) becomes a span inside a single root trace.

A typical trace hierarchy looks like:

rag_pipeline (root)
│
├── retrieve_documents
│   ├── embed_text
│   ├── index_documents (only once)
│   └── retrieve_documents
│
├── llm_completion
│
└── evaluate_rag_output
    ├── evaluate_relevancy
    └── evaluate_hallucination

This structure gives you:

Full-System Visibility

Every question generates a complete execution tree revealing:

• what happened
• where it happened
• how long it took
• what went wrong

End-to-End Metrics

• token usage
• retrieval scores
• latency per component
• evaluation metrics

Rich Debugging Context

Each span stores:

• input messages
• embeddings preview
• retrieved context
• generated outputs
• error details

Continuous Quality Monitoring

Your evaluation step logs:

• a relevancy score
• a hallucination risk
• a final pass-or-fail quality metric

Langfuse becomes the single pane of glass for understanding your RAG pipeline’s behavior, serving as the missing observability layer that transforms a working prototype into a production-ready system.

Would you like immediate access to 3,457 images curated and labeled with hand gestures to train, explore, and experiment with … for free? Head over to Roboflow and get a free account to grab these hand gesture images.

Project Setup

Before we write a single line of RAG logic, the foundation must be solid: a clean folder structure, repeatable environment setup, deterministic configuration, and a reliable inference and observability stack. This section walks you through the project layout, how to launch vLLM and Langfuse via Docker Compose, how to install retrieval dependencies (FAISS and SentenceTransformers), and how to configure all components using a single config.yaml file.

Folder Structure Walkthrough

Your project is organized for production clarity, where each subsystem (RAG, LLM, agent, evaluation, and infrastructure) is isolated in its own module.

project-root/
│
├── configs/
│   └── config.yaml                # Central config: LLM, embeddings, RAG, agent, eval, Langfuse
│
├── data/
│   └── sample_docs.txt            # Example inputs for retrieval
│
├── src/
│   ├── config.py                  # Config loader utilities
│   ├── llm_utils.py               # OpenAI-compatible client initialization
│   ├── llm_client.py              # Traced LLM wrapper (retry + token usage + spans)
│   ├── retriever.py               # FAISS retriever with traced indexing + search
│   ├── evaluation.py              # RAG quality scoring (relevancy + hallucination)
│   ├── rag_pipeline.py            # Full retrieval → prompt → generation → evaluation pipeline
│   ├── agent_orchestration.py     # 3-step traced agent workflow
│   ├── langfuse_instrumentation.py# Bootstraps Langfuse + flush utilities
│
├── docker-compose.yml             # vLLM + Langfuse + Postgres (self-hosted observability)
│
├── requirements.txt               # Python dependencies
│
└── check_rag_health.py            # Full system health check (env, docker, dependencies, files)

This layout ensures:

• Decoupled components: easy for testing and future replacement
• Reproducible environment: config-driven behavior
• Portable observability: one command launches everything
• Scalable structure: supports RAG, agents, and future tools

Every file in the src/ directory corresponds to a runnable pipeline step, and each is instrumented with Langfuse decorators so all activity becomes visible in the dashboard.

Starting vLLM and Langfuse Using Docker Compose

For production-like observability, the system relies on 2 running services:

• Langfuse: tracing, metrics, and span visualization
• vLLM: inference engine serving the LLM

Both are provided in your docker-compose.yml, and both run locally, meaning:

• zero cloud dependency
• zero per-token cost
• repeatable development environment

Start the entire stack (GPU version)

docker-compose --profile gpu up -d

Or start CPU mode

docker-compose --profile cpu up -d

Confirm services are running

docker-compose ps

You should see something like:

UI access

• Langfuse dashboard: http://localhost:3000
• vLLM API: http://localhost:8000/v1

What these services do internally

• Langfuse server: stores traces, spans, scoring, and metadata
• Langfuse worker: processes asynchronous scoring and analytics
• PostgreSQL: stores trace data
• vLLM: serves the Llama 2 model loaded at runtime

This cluster forms your local, production-grade observability and inference backbone.

Installing FAISS and SentenceTransformers

The retrieval layer requires:

• FAISS: similarity search
• SentenceTransformers: embedding model

These are already declared in your requirements.txt:

sentence-transformers>=2.2.0
faiss-cpu>=1.7.4
numpy>=1.24.0

Install dependencies:

pip install -r requirements.txt

After installation, verify FAISS is working:

python -c "import faiss; print(f'FAISS version: {faiss.__version__}')"

Verify embedding model loads:

python -c "from sentence_transformers import SentenceTransformer; print(SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2'))"

These 2 libraries form the core retrieval engine:

• Embeddings: produced locally (MiniLM)
• Retrieval: performed locally (FAISS IndexFlatL2)

No external API latency.

No cost.

No vendor lock-in.

Configuring config.yaml (LLM, Embeddings, RAG, and Evaluation)

The entire RAG and agent system is configurable from a single YAML file:

langfuse:
  host: "http://localhost:3000"
  project_name: "rag-selfhosted"

llm:
  base_url: "http://localhost:8000/v1"
  model: "meta-llama/Llama-2-7b-chat-hf"
  temperature: 0.7
  max_tokens: 300
  max_retries: 2

embeddings:
  model: "sentence-transformers/all-MiniLM-L6-v2"
  dimension: 384

rag:
  top_k: 3

agent:
  max_steps: 3

evaluation:
  enable_scoring: true
  min_quality_score: 0.6

Key configuration sections

LLM Configuration

Controls inference behavior:

• model name
• sampling temperature
• max tokens
• retry count
• endpoint (vLLM server URL)

Your TracedLLMClient loads these automatically via:

from config import get_llm_config

Embeddings Configuration

Controls vector dimension and embedding model, and is consumed by:

from config import get_embeddings_config

RAG Settings

Controls retrieval behavior:

• top_k: results returned from FAISS
• used inside: TracedRetriever.retrieve()

Agent Settings

Agent workflows build on top of RAG, controlling:

• max agent steps
• model used for intent detection

Evaluation Configuration

Defines quality control thresholds:

• relevancy
• hallucination risk
• minimum acceptable quality

Used in:

from config import get_evaluation_config

This config-driven system makes your pipeline:

• reproducible
• tunable
• production-friendly
• environment-agnostic

Change models or thresholds, and no code changes are required.

Need Help Configuring Your Development Environment?

All that said, are you:

• Short on time?
• Learning on your employer’s administratively locked system?
• Wanting to skip the hassle of fighting with the command line, package managers, and virtual environments?
• Ready to run the code immediately on your Windows, macOS, or Linux system?

Then join PyImageSearch University today!

Gain access to Jupyter Notebooks for this tutorial and other PyImageSearch guides pre-configured to run on Google Colab’s ecosystem right in your web browser! No installation required.

And best of all, these Jupyter Notebooks will run on Windows, macOS, and Linux!

Building a Langfuse-Traced Retriever with FAISS

The retriever is the beating heart of any RAG pipeline. If retrieval is weak, every downstream component (i.e., prompting, LLM generation, and evaluation) will degrade. In this section, we construct a production-grade retriever built on three pillars: local embeddings, FAISS vector search, and Langfuse instrumentation. The result is a component that is fast, reproducible, fully observable, and cheap to run because it never leaves your machine or calls a cloud API.

Your TracedRetriever class in src/retriever.py handles 4 responsibilities:

• Load an embedding model
• Embed and index documents
• Perform similarity-based search
• Emit Langfuse spans for every step (embedding, indexing, and retrieval)

Loading and Embedding Documents

The retriever begins by loading a local SentenceTransformers model, which provides dense vector embeddings without any external API calls.

embeddings_config = get_embeddings_config()
model_name = embeddings_config.get("model", "sentence-transformers/all-MiniLM-L6-v2")
self.model = SentenceTransformer(model_name)

Why local embeddings?

• No rate limits or API costs after the local environment is configured
• Fast inference through optimized ONNX or Torch acceleration
• Privacy-safe since no data leaves the environment
• Predictable latency, which is critical in production

Embedding a document (with tracing)

Your embed() method is wrapped with the Langfuse @observe decorator:

@observe(name="embed_text")
def embed(self, text: str) -> np.ndarray:

This automatically creates a top-level span called embed_text in Langfuse.

Inside the span, you record:

• first 100 characters of the text
• embedding dimension

langfuse_context.update_current_observation(
    input={"text_preview": text[:100]}
)

The embedding call itself:

embedding = self.model.encode([text], normalize_embeddings=True)[0]

This normalization step ensures the embeddings have unit length, which stabilizes similarity scoring and produces better retrieval in FAISS L2 spaces.

The span finishes by storing metadata:

langfuse_context.update_current_observation(
    output={"embedding_dim": len(embedding)}
)

This is extremely useful later for debugging:

• Did documents produce embeddings with inconsistent lengths?
• Are embeddings accidentally empty?
• Are overly long texts being passed in?

Langfuse gives you full visibility.

Creating and Populating a FAISS Index

After loading the embedding model, the retriever constructs a FAISS IndexFlatL2 index:

self.index = faiss.IndexFlatL2(self.dimension)

This index:

• Stores vectors in RAM
• Uses Euclidean distance (L2) for similarity
• Has no training step, making it ideal for small and medium-sized datasets

Your retriever keeps an in-memory list of source documents:

self.documents = []

Indexing documents (with tracing)

@observe(name="index_documents")
def index_documents(self, documents: List[str]):

This span tracks:

• how many documents are being indexed
• how many embeddings were added
• previews of content for debugging

Under the hood:

• Store the raw documents
• Embed them in a batch
• Add the vectors to FAISS

embeddings = self.model.encode(documents, normalize_embeddings=True)
self.index.add(embeddings.astype(np.float32))

Because FAISS expects float32, the cast is mandatory.

Why IndexFlatL2?

• Simple
• Deterministic
• Fast for small–medium corpora (< 200k docs)
• Plays well with normalized embeddings (MiniLM, BERT, etc.)

Your pipeline achieves high throughput without additional libraries or GPUs.

Retrieving with Similarity Ranking

The retrieval process begins with the retrieve() method:

@observe(name="retrieve_documents")
def retrieve(self, query: str, top_k: int = None):

Langfuse creates a tracing span named retrieve_documents for every search operation.

Step 1. Embed the query

query_embedding = self.embed(query).reshape(1, -1)

Notice that calling self.embed() creates a nested span under the retrieval span in Langfuse.

This nesting hierarchy:

retrieve_documents
    ├── embed_text

gives you a complete view of:

• how long embedding took
• token count (if embedding model changes)
• exact query text

Step 2. Search the FAISS index

distances, indices = self.index.search(query_embedding, top_k)

FAISS returns:

• indices: the closest documents
• distances: L2 distances to each doc

You convert distances into similarity scores:

relevance_score = 1.0 / (1.0 + float(distance))

This transforms smaller distances into higher scores.

Step 3. Format ranked results

results.append({
    "content": self.documents[idx],
    "score": relevance_score,
    "rank": rank + 1,
    "distance": float(distance)
})

Step 4. Log retrieval metadata to Langfuse

langfuse_context.update_current_observation(
    output={
        "result_count": len(results),
        "scores": [r["score"] for r in results],
        "results": [...]
    }
)

You even send content previews (200 characters), which appear in the Langfuse UI and make debugging dramatically easier.

Adding Langfuse Spans to Indexing and Retrieval Steps

Langfuse observability is woven into every retrieval path using the @observe decorator and metadata updates.

Spans you automatically get from your retriever

These spans appear under your RAG pipeline trace like:

rag_pipeline
    ├── retrieve_documents
    │       ├── embed_text
    │       └── result metadata
    ├── llm_completion
    ├── evaluate_rag_output
    └── final scoring

Why this matters in production

• You can identify whether latency is coming from embedding, FAISS search, or LLM inference.
• You can detect mismatches like:
  • wrong embedding dimension
  • missing documents
  • unnormalized vectors
  • misconfigured top-k
• You get complete end-to-end lineage for every query.
• You can monitor retriever performance across time.

This is the observability layer that most open-source RAG tutorials never include, but you now have it baked into the core of your retriever.

Building a Traced LLM Wrapper for vLLM and Langfuse

OpenAI-Compatible Chat Completions via vLLM

Your LLM wrapper is split into 2 layers:

• a low-level OpenAI-compatible client in llm_utils.py, and
• a high-level, Langfuse-traced wrapper in llm_client.py.

The low-level client is created in get_llm_client():

from openai import OpenAI

def get_llm_client(timeout: int = 60, load_model_from_config: bool = False):
    if os.getenv("OPENAI_BASE_URL") is None:
        print("⚠️  OPENAI_BASE_URL not found in environment. Using default http://localhost:8000/v1")
   
    if os.getenv("OPENAI_API_KEY") is None:
        print("⚠️  OPENAI_API_KEY not set. Using dummy key.")
   
    client = OpenAI(
        base_url=os.getenv("OPENAI_BASE_URL", "http://localhost:8000/v1"),
        api_key=os.getenv("OPENAI_API_KEY", "dummy"),
        timeout=timeout,
    )
    ...
    return client

This means that as long as vLLM is running behind an OpenAI-compatible server (from docker-compose.yml on http://localhost:8000/v1), the rest of your code simply calls client.chat.completions.create(...) exactly like it would against OpenAI, without any vendor-specific changes.

At the higher level, TracedLLMClient in src/llm_client.py wraps this client and pulls model configuration from configs/config.yaml via get_llm_config():

llm:
  base_url: "http://localhost:8000/v1"
  model: "meta-llama/Llama-2-7b-chat-hf"
  temperature: 0.7
  max_tokens: 300
  max_retries: 2

class TracedLLMClient:
    def __init__(self, model: str = None, max_retries: int = 2, timeout: int = 60):
        self.client = get_llm_client(timeout=timeout)
        if model is None:
            llm_config = get_llm_config()
            model = llm_config.get("model", "meta-llama/Llama-2-7b-chat-hf")
        self.model = model
        self.max_retries = max_retries

The end result: your RAG and agent code never talks to vLLM directly; it always goes through TracedLLMClient, which is OpenAI-compatible, config-driven, and ready for tracing.

Retry Logic and Error Handling

The core of the wrapper is the complete() method:

from langfuse.decorators import observe, langfuse_context

class TracedLLMClient:
    @observe(name="llm_completion")
    def complete(self, messages: List[Dict[str, str]], **kwargs) -> Dict:
        llm_config = get_llm_config()
        temperature = kwargs.get("temperature", llm_config.get("temperature", 0.7))
        max_tokens = kwargs.get("max_tokens", llm_config.get("max_tokens", 300))

        langfuse_context.update_current_observation(
            input={"messages": messages, "model": self.model}
        )

        last_error = None
        for attempt in range(self.max_retries):
            try:
                response = self.client.chat.completions.create(
                    model=self.model,
                    messages=messages,
                    temperature=temperature,
                    max_tokens=max_tokens
                )
                ...
                return {..., "success": True}
            except Exception as e:
                last_error = e
                if attempt < self.max_retries - 1:
                    time.sleep(1)
                    continue

        error_msg = f"LLM call failed after {self.max_retries} attempts: {last_error}"
        langfuse_context.update_current_observation(
            level="ERROR",
            output={"error": error_msg}
        )
        return {"content": None, "error": error_msg, "success": False}

A few production-grade details are baked in here:

• Config-driven defaults: temperature and max_tokens come from config.yaml but can be overridden per-call via kwargs.
• Retry loop: the method tries up to self.max_retries times (default 2), with a short time.sleep(1) backoff between attempts.
• Graceful failure: if all attempts fail, you get a structured response {content: None, error: "...", success: False} instead of a hard crash, and the Langfuse span is explicitly marked as "ERROR".

When you call this from rag_pipeline.py or agent_orchestration.py, you can safely check result["success"] and decide whether to return a fallback answer, propagate the error, or short-circuit the pipeline.

Logging Request and Response Payloads

Because TracedLLMClient is decorated with @observe(name="llm_completion"), every call automatically becomes a Langfuse span, and you manually enrich that span with inputs and outputs via langfuse_context.

At the start of the call, you log the request payload:

langfuse_context.update_current_observation(
    input={"messages": messages, "model": self.model}
)

This means that in the Langfuse UI you will see:

• the full chat history (messages) you sent to the model
• which model was used (e.g., "meta-llama/Llama-2-7b-chat-hf")

After a successful LLM call, you log the response content:

content = response.choices[0].message.content

langfuse_context.update_current_observation(
    output={"content": content},
    usage={
        "input": response.usage.prompt_tokens,
        "output": response.usage.completion_tokens,
        "total": response.usage.total_tokens
    },
    metadata={"attempt": attempt + 1}
)

So every Langfuse span for llm_completion will show:

• the raw answer text the model generated
• which attempt succeeded (first try or retry)
• the token usage for that call

On failure, the wrapper logs the error message instead of content:

langfuse_context.update_current_observation(
    level="ERROR",
    output={"error": error_msg}
)

This gives you debuggable traces when vLLM is down, you hit timeouts, or your model name is misconfigured.

Capturing Token Usage and Metadata in Langfuse

vLLM exposes response.usage in an OpenAI-like shape, and you forward that directly into Langfuse as part of the span:

langfuse_context.update_current_observation(
    output={"content": content},
    usage={
        "input": response.usage.prompt_tokens,
        "output": response.usage.completion_tokens,
        "total": response.usage.total_tokens
    },
    metadata={"attempt": attempt + 1}
)

return {
    "content": content,
    "usage": response.usage.model_dump(),
    "success": True
}

This gives you 2 layers of observability:

• Inside Langfuse
  • You can filter and inspect spans by usage.total, see which prompts are expensive, and spot unusually long generations.
  • You can correlate token usage with overall RAG or agent traces because llm_completion spans sit inside higher-level pipeline spans such as rag_pipeline or agent_workflow.
• Inside your Python code
  • Callers receive result["usage"] and can log or aggregate it themselves (e.g., cost dashboards, quotas, or alerting in future lessons).
  • Because usage is returned as response.usage.model_dump(), it is just a normal Python dict you can serialize or send elsewhere.

The metadata={"attempt": attempt + 1} block gives you a clean way to see how often retries are needed; if you start seeing a lot of second or third attempts in Langfuse, you know vLLM or your infra is becoming unreliable and needs attention.

Example: Using the Traced LLM Client

Your __main__ block in llm_client.py shows a minimal end-to-end example:

if __name__ == "__main__":
    client = TracedLLMClient()
   
    result = client.complete(
        messages=[
            {"role": "user", "content": "What is RAG in AI?"}
        ]
    )
   
    print(f"Response: {result['content']}")
    print(f"Tokens: {result['usage']['total_tokens']}")
   
    trace_id = langfuse_context.get_current_trace_id()
    langfuse_host = os.getenv("LANGFUSE_HOST", "http://localhost:3000")
    print(f"🔍 View trace: {langfuse_host}/trace/{trace_id}")

This script:

• verifies that vLLM is responding correctly
• verifies that Langfuse keys and host are properly configured
• gives you a direct URL to the exact trace for this LLM call in the Langfuse UI

In the next sections, you will see this same TracedLLMClient reused inside the RAG pipeline and RAG evaluation, where it becomes just one span in a larger, nested trace tree.

Building a Fully Traced RAG Pipeline with Langfuse

The run_rag_pipeline Orchestrator

Your full RAG flow is implemented in src/rag_pipeline.py as a single orchestrator function:

@observe(name="rag_pipeline")
def run_rag_pipeline(
    question: str,
    retriever: TracedRetriever,
    llm_client: TracedLLMClient,
    top_k: int = 3
) -> Dict:
    ...

This one function wires together everything you have built so far: it takes a user question, uses the traced retriever to find context, calls the traced LLM client to generate an answer, and then runs RAG evaluation to compute relevancy and hallucination scores. Because it is decorated with @observe(name="rag_pipeline"), the entire run shows up in Langfuse as a top-level trace, with all retrieval, LLM, and evaluation spans nested underneath.

Step 1: Retrieve

The first step is retrieving documents with your TracedRetriever:

langfuse_context.update_current_observation(
    input={"question": question, "top_k": top_k}
)

print("Step 1: Retrieving documents...")
docs = retriever.retrieve(query=question, top_k=top_k)

if not docs:
    print("❌ No documents found")
    return {"answer": "No relevant information found.", "success": False}

Here is what happens in this step:

• The pipeline span is enriched with the incoming question and the top_k parameter via langfuse_context.update_current_observation.
• retriever.retrieve(...) is itself decorated with @observe(name="retrieve_documents"), so Langfuse automatically creates a child span under rag_pipeline. Inside that span, you log the query, scores, and content previews.
• If the index is empty or nothing is returned, you fail fast with a friendly message and success=False instead of trying to prompt the LLM with no context.

By the end of Step 1, you have a ranked list of documents such as:

[
    {"content": "...", "score": 0.93, "rank": 1, "distance": 0.12},
    {"content": "...", "score": 0.88, "rank": 2, "distance": 0.18},
    ...
]

and their retrieval details are already captured in Langfuse.

Step 2: Build Prompt from Retrieved Docs

Next, you turn those retrieved documents into a single, structured prompt:

print("Step 2: Building prompt...")
context = "\n\n".join([f"[{i+1}] {d['content']}" for i, d in enumerate(docs)])
messages = [
    {"role": "system", "content": "Answer based on the provided context."},
    {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:"}
]

A few important details:

• Each document is tagged with an index ([1], [2], [3]) so it is easy to map parts of the final answer back to specific sources, both as a human and when you are debugging traces.
• The system message explicitly constrains the model: “Answer based on the provided context.” This is a simple but effective guardrail against hallucinations.
• The user message includes both the stitched context and the original question, finishing with “Answer:” to bias the model toward a direct response.

Because messages are later passed into the traced LLM client, the entire prompt (including context) is visible inside the llm_completion span in Langfuse.

Step 3: Generate with vLLM

You then hand off the prompt to your TracedLLMClient:

print("Step 3: Generating answer...")
result = llm_client.complete(messages)

if not result["success"]:
    print(f"❌ Generation failed: {result.get('error')}")
    return {"answer": None, "error": result.get("error"), "success": False}

answer = result["content"]
print(f"✅ Answer generated\n")

Under the hood:

• TracedLLMClient.complete() calls client.chat.completions.create(...) against the vLLM OpenAI-compatible server (configured via OPENAI_BASE_URL and OPENAI_API_KEY, with model and temperature from config.yaml).
• The method is decorated with @observe(name="llm_completion"), so a child span is created inside the rag_pipeline trace.
• Inside that span, you log:
  • the messages and model as input
  • the generated content as output
  • detailed token usage (prompt_tokens, completion_tokens, total_tokens) as usage, plus metadata={"attempt": ...} indicating which retry succeeded

If vLLM is down, misconfigured, or times out, the wrapper returns {"success": False, "error": ...} and updates the Langfuse span with level="ERROR", so you get a clear red node in the trace instead of a mysterious failure.

Step 4: Evaluate Response Quality

Once you have an answer, you pass everything into the evaluation layer in src/evaluation.py:

print("Step 4: Evaluating quality...")
evaluation_results = evaluate_rag_output(question, docs, answer)

evaluate_rag_output is itself annotated with @observe(name="evaluate_rag_output"), and it calls 2 more traced helpers under the hood:

• evaluate_relevancy(query, retrieved_docs, answer)
• evaluate_hallucination_risk(retrieved_docs, answer)

The process inside evaluate_rag_output looks like this:

langfuse_context.update_current_observation(
    input={
        "query": query,
        "doc_count": len(retrieved_docs),
        "answer_length": len(answer)
    }
)

relevancy_score = evaluate_relevancy(query, retrieved_docs, answer)
hallucination_risk = evaluate_hallucination_risk(retrieved_docs, answer)
overall_quality = (relevancy_score + (1.0 - hallucination_risk)) / 2.0

eval_config = get_evaluation_config()
min_quality = eval_config.get("min_quality_score", 0.6)

results = {
    "relevancy_score": relevancy_score,
    "hallucination_risk": hallucination_risk,
    "overall_quality": overall_quality,
    "passed": overall_quality >= min_quality
}

In more detail:

• Relevancy scoring (evaluate_relevancy): computes how well the answer overlaps with both the query and the retrieved documents using simple word-level heuristics.
• Hallucination risk (evaluate_hallucination_risk): estimates how many of the answer’s content words are grounded in the retrieved documents; low grounding means higher risk.
• Overall quality: is a simple average of relevancy and 1 − hallucination_risk, giving a single number between 0 and 1.
• A minimum quality threshold (min_quality_score) comes from the evaluation section of config.yaml and is used to set a passed boolean.

The function then:

langfuse_context.score_current_observation(
    name="relevancy",
    value=relevancy_score,
    comment="Keyword and document relevance"
)

langfuse_context.score_current_observation(
    name="hallucination_risk",
    value=hallucination_risk,
    comment="Risk of ungrounded claims"
)

langfuse_context.score_current_observation(
    name="overall_quality",
    value=overall_quality,
    comment=f"Combined quality score (threshold: {min_quality})"
)

langfuse_context.update_current_observation(output=results)

So you get 3 named scores on the evaluation span inside Langfuse: relevancy, hallucination_risk, and overall_quality, each with a numeric value and a human-readable comment.

Tracing the Entire RAG Pipeline with Nested Spans

Back in run_rag_pipeline, you finalize the top-level observation and return a structured result:

langfuse_context.update_current_observation(
    output={
        "answer": answer,
        "sources_count": len(docs),
        "evaluation": evaluation_results
    }
)

print(f"✅ Evaluation complete")
print(f"  Relevancy: {evaluation_results['relevancy_score']:.2f}")
print(f"  Hallucination Risk: {evaluation_results['hallucination_risk']:.2f}")
print(f"  Overall Quality: {evaluation_results['overall_quality']:.2f}")
print(f"  Passed: {'✅' if evaluation_results['passed'] else '❌'}\n")

Then you expose the trace URL:

trace_id = langfuse_context.get_current_trace_id()
langfuse_host = os.getenv("LANGFUSE_HOST", "http://localhost:3000")

print(f"{'='*50}")
print(f"✅ Pipeline Complete")
print(f"🔍 View trace: {langfuse_host}/trace/{trace_id}")
print(f"{'='*50}\n")

At this point, a single pipeline run creates a hierarchy of spans roughly like this in Langfuse:

• rag_pipeline (top-level)
  • retrieve_documents (from TracedRetriever.retrieve)
    • embed_text (from TracedRetriever.embed)
  • llm_completion (from TracedLLMClient.complete)
  • evaluate_rag_output
    • evaluate_relevancy
    • evaluate_hallucination

Each node contains its own inputs, outputs, usage, and scores, giving you a complete picture of where time is spent, how the model behaved, and whether the final answer passed your quality threshold.

Returned Structure and Downstream Use

Finally, the function returns a rich Python dictionary :

return {
    "answer": answer,
    "sources": docs,
    "evaluation": evaluation_results,
    "success": True
}

This shape is deliberate:

• answer: can be rendered in a UI, CLI, or logged for later inspection.
• sources: lets you show which documents backed the answer (e.g., for “source citations” in a frontend).
• evaluation: gives your downstream systems a simple way to gate responses (e.g., only show answers where overall_quality >= 0.7).
• success: makes it easy to distinguish between “no documents”, “LLM error”, and “normal completion”.

Together, this section gives you not just a RAG pipeline, but a fully traced, quality-scored RAG system that is ready to plug into dashboards, UIs, or further production hardening.

Implementing LLM Evaluation Metrics for RAG: Relevancy and Hallucination Risk

Relevancy Scoring

Relevancy is implemented in evaluate_relevancy() and answers a simple but crucial question: “How well does the model’s answer align with the retrieved documents and the user’s query?”

Your scoring function uses a lightweight, keyword-overlap heuristic, which is ideal for debugging and observability without introducing another model dependency. The implementation:

@observe(name="evaluate_relevancy")
def evaluate_relevancy(query: str, retrieved_docs: List[Dict], answer: str) -> float:
    langfuse_context.update_current_observation(
        input={"query": query, "doc_count": len(retrieved_docs), "answer_length": len(answer)}
    )

    query_words = set(query.lower().split())
    answer_words = set(answer.lower().split())

    overlap_with_query = len(answer_words & query_words) / max(len(query_words), 1)

    doc_words = set()
    for doc in retrieved_docs:
        doc_words |= set(doc["content"].lower().split())

    overlap_with_docs = len(answer_words & doc_words) / max(len(answer_words), 1)

    relevancy_score = (overlap_with_query + overlap_with_docs) / 2.0

    langfuse_context.score_current_observation(
        name="relevancy",
        value=relevancy_score,
        comment="Keyword and doc overlap relevance"
    )

    langfuse_context.update_current_observation(output={"relevancy": relevancy_score})
    return relevancy_score

What the algorithm evaluates:

• Query–Answer overlap: Ensures the model is addressing the question.
• Document–Answer overlap: Checks that the model grounds its answer in retrieved context.
• The final score is the average of both signals.

While simple, this gives you an interpretable, production-friendly metric that appears directly in Langfuse traces.

Hallucination Risk Estimation

Hallucination risk is implemented in evaluate_hallucination_risk() and estimates how much of the answer is unsupported by the retrieved documents.

@observe(name="evaluate_hallucination")
def evaluate_hallucination_risk(retrieved_docs: List[Dict], answer: str) -> float:
    all_doc_words = set()
    for doc in retrieved_docs:
        all_doc_words |= set(doc["content"].lower().split())

    answer_words = set(answer.lower().split())

    grounding_ratio = len(answer_words & all_doc_words) / max(len(answer_words), 1)
    hallucination_risk = 1.0 - grounding_ratio

Interpretation:

• If every important token in the answer appears in the retrieved context, the hallucination risk is low.
• If the answer relies heavily on tokens not present in any source document, the hallucination risk is high.

Langfuse logs this as:

langfuse_context.score_current_observation(
    name="hallucination_risk",
    value=hallucination_risk,
    comment="Ungrounded token ratio"
)

This trace node helps you immediately visualize how close an answer is to going “off the rails.”

Overall Quality Metric

Your master scoring function evaluate_rag_output() combines the 2 metrics:

overall_quality = (relevancy_score + (1.0 - hallucination_risk)) / 2.0

This means:

• high relevancy and low hallucination risk indicate high quality
• low relevancy and high hallucination risk indicate low quality

config.yaml defines the minimum acceptable score:

evaluation:
  min_quality_score: 0.6

Then:

passed = overall_quality >= min_quality

This allows your downstream systems to treat RAG evaluation like a gatekeeper:

• passed=True: show the answer to the user, store it, or send it downstream
• passed=False: trigger fallback mode, self-reflection, or agentic repair workflows

All 3 metrics (relevancy, hallucination_risk, and overall_quality) are scored and attached to the current Langfuse span.

How Langfuse Displays Evaluation and Scoring Nodes

The evaluation subsystem produces one of the most informative trace segments in Langfuse. A typical structure:

rag_pipeline
 ├── retrieve_documents
 ├── llm_completion
 └── evaluate_rag_output
       ├── relevancy (score)
       ├── hallucination_risk (score)
       ├── overall_quality (score)

Each node includes:

Inputs

• user query
• document count
• answer length

Outputs

• numeric scores
• pass-or-fail status
• evaluation metadata

Visual Benefits Inside Langfuse

• Color-coded score nodes help you spot failing RAG runs instantly.
• Timeline alignment shows you evaluation overhead and where bottlenecks appear.
• Nested spans reveal exactly which part of the pipeline caused a failure.
• JSON detail view allows exporting evaluation metrics for dashboards or analytics.

With these evaluation spans, your Langfuse trace evolves from a simple log viewer into a quality monitoring dashboard for your RAG system.

Running and Inspecting the RAG Pipeline End-to-End

Running rag_pipeline.py End-to-End

With all components in place (the retriever, the traced LLM wrapper, and the evaluation module), you can now run the complete production-grade RAG pipeline. The script rag_pipeline.py orchestrates the entire flow:

python src/rag_pipeline.py

This script loads documents, indexes them, retrieves the top_k matches, builds a contextual prompt, generates an answer using vLLM, evaluates the output quality, and logs every step into Langfuse. If all services are running (Langfuse UI on port 3000 and vLLM on port 8000), the run completes with a final console message showing the trace URL:

🔍 View trace: http://localhost:3000/trace/<trace_id>

This makes it trivial to jump directly into the corresponding trace in your observability dashboard and inspect the entire RAG execution, including nested spans and evaluation scores.

Example Trace Outputs

A successful run produces a hierarchical trace structure in Langfuse that mirrors your pipeline architecture. A typical RAG trace looks like this:

rag_pipeline
 ├── retrieve_documents
 │     ├── embed_text
 │     └── FAISS search metadata
 ├── llm_completion
 │     ├── request payload
 │     ├── response payload
 │     └── token usage
 └── evaluate_rag_output
        ├── relevancy (score)
        ├── hallucination_risk (score)
        └── overall_quality (score)

What you will see in the trace:

Retrieval metadata

• top_k value
• query text
• relevance scores
• FAISS distances
• document preview snippets

LLM generation metadata

• system and user messages used for prompting
• token usage breakdown
• retry attempts
• vLLM latency and response time

Evaluation metrics

• numeric relevancy score
• hallucination risk estimation
• overall quality score
• pass-or-fail decision using the threshold in config.yaml

Together, these give you a full audit trail for each RAG run, which is perfect for debugging, monitoring, or offline analysis.

Debugging with the Langfuse UI (Span Trees, Scores, and Metadata)

Langfuse is not just a logger; it acts as a visual debugger for your entire RAG system. When you open the trace URL, you will see several powerful debugging tools:

Span Tree View

This hierarchical tree shows the exact execution order and timing of:

• retrieval
• embedding
• indexing
• LLM generation
• evaluation steps

It helps you detect:

• slow spans (bottlenecks)
• failed or retried LLM calls
• missing or empty retrieval results

Scoring Nodes

Evaluation scores appear as structured nodes:

• relevancy
• hallucination_risk
• overall_quality

Langfuse color-codes these (green, yellow, and red), making it instantly clear when a RAG answer is degrading in quality.

Metadata Panels

Each span contains:

• input and output payloads
• token counts
• FAISS distances
• processed document counts
• retry counts
• trace-level summaries

This makes debugging extremely fast:

• Wrong documents retrieved? Inspect retrieval span input and output.
• Unexpected LLM answer? Check the exact prompt in the generation span.
• Poor evaluation scores? Expand the scoring spans to see the raw metrics.

Because traces are stored locally in your self-hosted Langfuse instance, you get complete transparency without relying on cloud telemetry.

Viewing RAG Traces, Spans, and Scores in Langfuse

Once your RAG pipeline is running end-to-end, the real magic happens inside Langfuse. This is where retrieval steps, LLM calls, evaluation metrics, token usage, and pipeline-level metadata condense into a single, navigable trace. In this section, you will learn how to interpret that trace, span by span, so you can debug, understand, and improve RAG behavior with production-grade visibility.

Understanding Hierarchical Spans (Retrieve → Prompt → Generate → Evaluate)

Langfuse automatically groups each step of your rag_pipeline into a nested hierarchy of spans. A typical RAG trace looks like this:

• rag_pipeline (root trace)
  • retrieve_documents
    • embed_text
  • llm_completion
  • evaluate_rag_output
    • evaluate_relevancy
    • evaluate_hallucination
  • scoring nodes (overall_quality, relevancy, hallucination_risk)

This hierarchy corresponds directly to your source code:

• TracedRetriever.retrieve(): retrieval span
• TracedLLMClient.complete(): generation span
• evaluate_rag_output(): evaluation span

How to Navigate the Span Tree

Each span reveals:

• execution time (critical for latency bottlenecks)
• inputs and outputs captured via langfuse_context.update_current_observation()
• whether nested operations (e.g., embedding calls) executed successfully
• metadata from FAISS search, document previews, and query text

Langfuse becomes a timeline and debugger for your RAG system.

Inspecting Retrieval: Document Scores and Previews

The retrieval stage is your first major insight point. The retrieve_documents span logs:

• the query that was embedded
• the top_k used for FAISS search
• distance scores returned
• converted relevancy scores (your 1/(1+d) heuristic)
• ranked documents with text previews

What to Look For

• High distances and low scores: embedding mismatch or poor docs
• Same document repeatedly ranking #1: indexing error
• Empty results: index not built or FAISS dimension mismatch

Embedded Text Span

The embed_text span reveals the preview of text used for embeddings:

• inspect embeddings length
• detect empty or malformed documents
• verify embeddings model configuration

Inspecting Prompt Construction (Optional View)

Prompt creation happens between retrieval and generation. Although you do not create a separate Langfuse span for this step, the constructed prompt appears inside the LLM span input.

What you verify here:

• context formatting
• numbering
• whitespace
• hallucination-reducing systems instructions

This becomes essential when debugging wrong answers.

Token Usage and Generation Metadata

Inside the llm_completion span, Langfuse records:

• input tokens
• output tokens
• total tokens
• retry count
• model name
• latency breakdown
• response content

What to Look For

• Unusually high input tokens: prompt too large
• High output tokens: model drifting or verbose
• Repeated retries: vLLM throughput issue
• Very long latency: GPU under-provisioned or CPU fallback

Evaluation Scoring Nodes (Relevancy, Hallucination, and Overall Quality)

Your evaluation functions (evaluate_relevancy, evaluate_hallucination_risk, evaluate_rag_output) create 3 scoring nodes inside Langfuse:

• relevancy
• hallucination_risk
• overall_quality

These appear alongside the evaluation span.

How to Interpret Them

• High relevancy and low hallucination risk: high overall_quality
• Low relevancy and high hallucination risk: RAG failure
• passed=True means the response met the min_quality_score threshold in config.yaml

Debugging Failures

• Relevancy low: retrieval needs improvement
• Hallucination high: prompt needs grounding
• Both low: LLM ignoring context, bad retrieval, or noisy docs

Visual Timeline and Performance Profiling

The timeline view shows exact timings:

• embedding
• retrieval
• prompt construction
• LLM generation
• evaluation

This allows profiling end-to-end latency.

How Langfuse Helps Production Debugging

Langfuse tracing helps answer real production questions:

“Why was this answer wrong?”

Open the evaluation spans, review the hallucination score, inspect the prompt, and then inspect the retrieved documents.

“Which part is slowing down?”

Open the timeline and locate the bottleneck, which is often embeddings or the LLM.

“Did the LLM actually use the retrieved documents?”

Compare:

• retrieval previews
• answer keywords
• relevancy score

“Why did this query fail?”

The trace will show:

• empty index
• retries
• exceptions
• malformed inputs
• missing environment variables

In production, this becomes indispensable.

What's next? We recommend PyImageSearch University.

Course information:
120+ total classes • 115+ hours hours of on-demand code walkthrough videos • Last updated: August 2026
★★★★★ 4.84 (128 Ratings) • 16,000+ Students Enrolled

I strongly believe that if you had the right teacher you could master computer vision and deep learning.

Do you think learning computer vision and deep learning has to be time-consuming, overwhelming, and complicated? Or has to involve complex mathematics and equations? Or requires a degree in computer science?

That’s not the case.

All you need to master computer vision and deep learning is for someone to explain things to you in simple, intuitive terms. And that’s exactly what I do. My mission is to change education and how complex Artificial Intelligence topics are taught.

If you're serious about learning computer vision, your next stop should be PyImageSearch University, the most comprehensive computer vision, deep learning, and OpenCV course online today. Here you’ll learn how to successfully and confidently apply computer vision to your work, research, and projects. Join me in computer vision mastery.

Inside PyImageSearch University you'll find:

• ✓ 120+ courses on essential computer vision, deep learning, and OpenCV topics
• ✓ 94+ Certificates of Completion
• ✓ 115+ hours hours of on-demand video
• ✓ Brand new courses released regularly, ensuring you can keep up with state-of-the-art techniques
• ✓ Pre-configured Jupyter Notebooks in Google Colab
• ✓ Run all code examples in your web browser — works on Windows, macOS, and Linux (no dev environment configuration required!)
• ✓ Access to centralized code repos for all 540+ tutorials on PyImageSearch
• ✓ Easy one-click downloads for code, datasets, pre-trained models, etc.
• ✓ Access on mobile, laptop, desktop, etc.

Click here to join PyImageSearch University

Summary

In this lesson, you built a fully instrumented, production-grade RAG pipeline and learned how observability transforms retrieval-augmented systems from “black boxes” into transparent, debuggable, measurable workflows. You started by setting up the core infrastructure (self-hosted Langfuse, vLLM for fast local inference, and FAISS and SentenceTransformers for efficient retrieval) and then wired all these components together using a clean, traceable architecture.

With tracing enabled end-to-end, every stage of your RAG pipeline became inspectable: document embedding, FAISS indexing, retrieval scoring, prompt construction, LLM generation, and quality evaluation. You saw how Langfuse automatically visualizes these steps as nested spans, how it captures token usage and metadata for LLM calls, and how your evaluation functions produce relevancy, hallucination risk, and overall-quality scores directly inside the trace.

By running the pipeline and examining the traces, you learned how to debug retrieval quality, diagnose prompt-related issues, inspect model behavior, and identify performance bottlenecks using Langfuse’s hierarchical tree view and timeline profiler. The final result is an observability-first RAG stack: fully local, fast, and transparent, designed exactly the way production systems must operate.

This foundation prepares you for upcoming lessons, where we extend the same tracing principles to multi-step agents, adding reasoning chains, intent analysis, and multi-span agent workflows on top of the RAG engine you constructed here.

Citation Information

Singh, V. “RAG Observability with Langfuse, vLLM, and FAISS,” PyImageSearch, S. Huot, A. Sharma, and P. Thakur, eds., 2026, https://pyimg.co/g20yk

@incollection{Singh_2026_rag-observability-langfuse-vllm-faiss,
  author = {Vikram Singh},
  title = {{RAG Observability with Langfuse, vLLM, and FAISS}},
  booktitle = {PyImageSearch},
  editor = {Susan Huot and Aditya Sharma and Piyush Thakur},
  year = {2026},
  url = {https://pyimg.co/g20yk},
}

To download the source code to this post (and be notified when future tutorials are published here on PyImageSearch), simply enter your email address in the form below!

Download the Source Code and FREE 17-page Resource Guide

Enter your email address below to get a .zip of the code and a FREE 17-page Resource Guide on Computer Vision, OpenCV, and Deep Learning. Inside you'll find my hand-picked tutorials, books, courses, and libraries to help you master CV and DL!

Download the code!

Website

The post RAG Observability with Langfuse, vLLM, and FAISS appeared first on PyImageSearch.