Introduction

In a recent project, I compared output from HuggingFace’s Candle with PyTorch equivalents to make sure that embeddings and calculations involving them are behaving as expected.

This led to comparing Python snippets that look like this:

def normalize_l2(embeddings):
    return F.normalize(embeddings, p=2, dim=1)

to Rust that looks like this:

pub fn normalize_l2(v: &Tensor) -> Result<Tensor> {
    Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?)
}

You might think this isn’t a fair comparison, because the normalization is hand-rolled in Rust and just tucked away in a function call in the Python equivalent. But that’s besides the point; even a hand-rolled Python equivalent would not be this noisy. PyTorch semantics are fraught with unintuitive behavior behind primitive operators. As a matter of fact, I wrote an earlier post about getting confused by PyTorch broadcasting mechanics when trying to normalize a tensor. Can you spot the problem?

P = N.float()
P /= P.sum(1)

Favoring Explicitness over Implicitness

When I was learning PyTorch (and I suppose NumPy, by extension), I found that I would often struggle to understand code that I read, because there were a lot of implicit interactions going on that I wasn’t aware of:

While these make for clean-looking code, it creates issues:

• Failures often occur silently and cause calculations to have incorrect results, sometimes subtly. This is worst kind of failure that you can have in software.
• There’s a notable absence of a type system that could catch logic failures and serve as documentation.

When I was learning, I found myself wishing that these operations were explicitly spelled out and enforced by a compiler. Without this, new developers must rely on someone (or more commonly, on documentation of inconsistent quality) to explain behavior that is not always very intuitive. It is always an option to determine this behavior by reading the source code, but unfortunately large projects such as PyTorch become extremely cumbersome to navigate once they grow past a certain size and complexity.

Forcing the Issues

The parts of Rust’s syntax which I used to perceive as noise. Let me paste in normalize_l2 again:

pub fn normalize_l2(v: &Tensor) -> Result<Tensor> {
    Ok(v.broadcast_div(&v.sqr()?.sum_keepdim(1)?.sqrt()?)?)
}

My early reaction to Rust code was to find the question marks, Result wrapping of the return type, and “OK” to be a lot of unfamiliar visual noise. These are some of are Rust’s way of forcing the programmer to acknowledge the potential failure points in their code. For a single developer who is continuously working on a codebase, you could argue this is not a big deal. Similarly, you could argue that it would create a burden for a small team of like-minded and extremely talented developers that is more trouble than its worth.

But the value quickly compounds when you consider other scenarios. If you ever want to onboard more junior developers, it is immensely helpful for them to be forced to understand where landmines are in software. If you walk away from the codebase for some period of time, or software development entirely, and some of those details will slip. Having a strict compiler enforcing hygeine is immensely beneficial for compensating for human error. The code is noisy because the task at hand is noisy.

Still, it would be nice if Rust could maintain its strictness with cleaner syntax. FP nerds probably recognize that the function is moving the input tensor through a Monadic context, and if the library existed in Haskell, it would be able to drop some of the noise with it’s “do notation”, like this:

normalizeL2 :: Tensor -> Either Error Tensor
normalizeL2 v = do
  squared <- sqr v
  summed <- sumKeepDim 1 squared
  norms <- sqrt summed
  broadcastDiv v norms

Rust adopted many of ML-style language’s ideas, but I don’t think anyone has figured out how to let it take Haskell’s “do” without giving up imperative control flow like FP languages often do, which would be at odds with it goal of being a systems programming language. If you are forced to give up either explicitness, speed, or elegance, I think Rust makes a good compromise.

Thanks to Matt for turning me into “that guy that keeps bringing Rust into the conversation”

Introduction

Reproducing GPT-2 from scratch is a great way to build an understanding of deep learning and large language models. A number of learning resources about GPT-2 architecture are available, including Andrej Karpathy’s free Zero to Hero course and Sebastian Raschka’s book Build a Large Language Model from Scratch. While these resources provide excellent coverage of how GPT-2 works under the hood, there’s a few details about challenges you’ll face in the wild when actually training on expensive hardware that aren’t extensively covered.

Understanding Cloud GPU Costs

It’s important to understand the process of provisioning servers and the details of what you’re doing to avoid racking up an unnecessarily large bill. I recently used a modified version of the reference pytorch implementation from Karpathy’s llm.c project to do my own reproduction of the model (This version appeared to be more up to date than NanoGPT and build-nanogpt variants at the time I ran it).

Many cloud providers offer VPSs with GPUs attached. I’ve seen Lambda Labs recommended here and there, so I checked their service first. One thing that I made sure to do was to ask for some trial credits - if they weren’t available, I would have chosen a different service. This instinct paid off because I made a number of mistakes and would have felt bad about paying the bill afterwards.

Lambda Labs instances worked with no fuss. If you’ve used VPSs like Digital Ocean’s droplets or Linode, Lambda’s instances work the same way. You upload your SSH public key and remote in to a ubuntu instance.

If like me, you hadn’t rented expensive GPUs before, it’s important to know how they are billed. I’m used to small linux VPSs that charge $5-$10/mo in the past, but servers set up for GPU workloads run for far more than that - about $24/hr on the instance I was interested in. Billing is a flat rate per hour, and it accumulates by the minute. It doesn’t matter whether you’re actively running any jobs, or even if the instance is currently running. The second you provision it the clock starts ticking, and you continue being billed until you destroy the instance and any data on it. This has major implications on how you’ll want to do your training.

Recommended Approach

Let’s say you wanted to train the network on 8 x H100s. If you rented a server with this setup, which costs about $24/hr, you would soon discover that you spend a good chunk of your time on data preparation. It involves downloading, tokenizing, and sharding the data - which in my case was the 10B sample of fineweb-edu. This is an IO bound operation, so all those expensive GPUs are not going to make it go any faster. If you run all that prep and your training run fails, you’re out $20.

My recommendation would be to start out with the cheapest server that is close to what you want to use for a full run. For example, if you wanted to eventually run on 8 X H100s, you would ideally want a server with multiple GPUs with at least 40GB of memory. That way you can fiddle with distributing loads and hyperparameters to improve your chances of getting things right on a real run.

My Run

After going through Karpathy’s course and skimming Raschka’s book, I decided to take a stab at reproducing GPT-2 based on some of the reference code from llm.c, with a few tweaks.

Data Preparation Phase

What I ended up doing was renting 4x 48GB A6000s for cheap and then performing the data preparation. As soon as the data preparation kicked off, I started compressing a copy of the project + the sharded data with tar -czvf llm.c.tar.gz llm.c. While that was working, I opened an additional client and started running some experiments, sanity checks, and small runs to make sure everything was working as expected. This gave me a chance to work out a few kinks that I didn’t know about.

After tar completed, I copied the archive back to my machine with scp -r -i ~/.ssh/id_ed25519 ubuntu@{Machine IP}:/home/ubuntu/llm.c/pylog_gpt2_124M .. Once that finished, I was done, so I shut down the instance and destroyed it.

Training Phase

Next, I tried spinning up an 8 x 80GB H100 instance. I moved my project and sharded data over with scp -C -i ~/.ssh/id_ed25519 llm.c.tar.gz ubuntu@{Machine IP}:/home/ubuntu/, and un-tar’d it with tar -xzvf llm.c.tar.gz, which took about half the time as the initial data preparation process. It probably would have been faster to upload it to a S3 bucket or something, but I don’t have one.

I tried running my code with a set of hyperparameters that I picked out, and the process failed with an OOM error. I was trying an aggressive mini-batch size, and unfortunately I didn’t feel comfortable tweaking the parameters to try to get it working since the bill racks up quickly. Instead, I blew up the instance to try different hyperparameters on a different set of hardware.

I had a hyperparameter set that I was sure would work with 8 x 40GB A1000s, so I spun that up and repeated the process. This time, it ran successfully to completion. I added some code that saved a model and did a little inferencing, so after that finished I copied the logs and model back over to my home machine and blew up this instance as well.

Conclusion

I’ve been slowly building up to understanding GPT-2’s architecture over the last few months and it was very fun to try out a training run. I can easily recommend the two learning resources mentioned above. If you have any questions about the process or get stuck, feel free to give me a holler.

Introduction

In a previous post, I explored using LLMs to perform text classification tasks. The idea was that they could enable automation of more complex tasks that are otherwise not automatable. Many people just use a model like ChatGPT for this, but after learning a bit about how they work, I was getting the sense that other approaches might be more efficient and accurate.

My experimental task was to classify messages sent from customers to an online store. Messages were classified as either refund requests, order inquiries, or general feedback.

In that post, I used a wrapper around llama.cpp, running inference on the 8B parameter, 4-bit quantized version of Llama 3. Architecturally, this is works like a tiny ChatGPT, and is capable of running on pretty much any consumer hardware. By directly prompting the base model I achieved ~92% accuracy out of the box, improving to ~95% with post-processing (retrying invalid responses).

The results seemed okay, but a general purpose instruction model didn’t really seem ideal for a specific task like this. I wanted to explore other options that might be more efficient or have better performance.

Two possible approaches to creating a task-specific model would be to either fine-tune an existing model or to train a smaller, specialized model from scratch.

Training a model from scratch would be extremely interesting, but it would also require a huge amount of data and computing power. On the other hand, fine-tuning is a technique that adapts pre-trained language models to specific tasks by further training on a smaller, task-specific dataset. This allows you to leverage the data and compute that went into the original pre-training of the model.

Model selection

My first naive idea was to try to fine-tune Llama 3 itself. Models such as GPT or Llama are primarily designed for text generation. Mechanically, they are trained to predict the next token in a sequence using masked self-attention, where each token can only attend to context to its left. What this means is that the neural network is only considering the words that appear before the token it is trying to predict. Lets look at an example.

Consider the sentence Live free or die hard. GPT/Llama will take this whole sentence and simultaneously learn to predict the next word given the input:

This makes perfect sense if you’re training a model to generate text, but it doesn’t really align with what we’re trying to do. We’re more interested in looking at an entire chunk of text and classifying it. It turns out there’s better ways to achieve this.

Other models learn other types of tasks, such as Masked Language Modeling and Next Sentence Prediction. In Masked Language Modeling, the model is tasked with predicting the missing words in a sentence. For example, given the sentence “The [MASK] jumped over the fence,” the model would try to predict the masked word (likely “dog” or “cat”).

In Next Sentence Prediction, the model is tasked with determining if two sentences are adjacent in a text. For instance, given the sentences “I love ice cream” and “It’s my favorite dessert”, the model would determine that these sentences are relatively likely to be adjacent. For “The sky is blue” and “Elephants have trunks”, the model would determine that these sentences relatively unlikely to be adjacent.

These approaches can be implemented through a mechanism called bidirectional self-attention. Unlike with masked-self attention, the model does not restrict itself to only consider tokens before something it is trying to predict. Instead, it considers all the tokens in the input and takes a stab at its task. This makes more sense if you’re trying to perform sentiment analysis, named entity recognition, or text classification. In those cases, you wouldn’t want all of the machinery associated with next-token prediction, because it would be dead weight.

BERT

A model that uses bidirectional self-attention is Bidirectional Encoder Representations from Transformers (BERT). Introduced by Google in 2018, BERT is pre-trained on Masked Language Modeling and Next Sentence Prediction described above.

These pre-training tasks equip BERT with a strong grasp of language structure and context. From the words of the authors, …the pre-trained BERT model can be fine-tuned with just one additional output layer to create state-of-the-art models for a wide range of tasks, such as question answering and language inference, without substantial task-specific architecture modifications.

My guess was that a model like this would be a better starting point for my use case.

DistilBERT

In 2019, Hugging Face released a smaller, distilled version of BERT called DistilBERT. It preserves 95% of BERT’s performance while using 40% fewer parameters and running 60% faster. For comparison, it is less than 1% of the size of Llama 3 8B.

DistilBERT achieves this efficiency through a process called knowledge distillation, where a smaller model (the student) is trained to mimic the behavior of a larger model (the teacher). This results in a more compact model that can still perform well on a wide range of tasks, making it well-suited for fine-tuning on specific applications, especially when computational resources are limited.

DistilBERT Architecture Overview

DistilBERT is based on the Transformer architecture. The model consists of several key components that work together to process and understand text. Text input is first passed through an embedding layer. This layer is responsible for converting input tokens (which are essentially words or parts of words) into dense vector representations. These vectors contain floating point numbers that encode the semantic meaning of the words in a format that the model can work with. The values of the vectors change throughout training as the model iterates on its semantic understanding.

The embedding layer learns that words like “dog” and “cat” are semantically closer to each other than to words like “car” or “building”. It captures subtle relationships, such as understanding that “king” is to “man” as “queen” is to “woman”, or that certain adjectives tend to precede certain types of nouns (like “delicious” often appearing before food-related words).

Once the text is embedded, it passes through multiple Transformer blocks. These blocks contain layers of self-attention and feed-forward neural networks. The self-attention mechanism allows the model to weigh the importance of different words in the input. This means the model can understand context and relationships between words, much like how we understand language by considering words in relation to each other. After the self-attention layer, feed-forward neural networks further process the attention output, allowing the model to capture complex patterns in the data.

Consider the sentence “The man who crossed the street was hit by a car.” Self-attention allows the model to understand that “hit” is more strongly related to both “man” and “car” than to “crossed” or “street”. This helps the model correctly interpret who was hit (the man) and what hit him (the car), even though these words are not adjacent in the sentence. The feed-forward networks then process this contextual information, potentially learning higher-level concepts like “traffic accidents” or “pedestrian safety” from such examples.

These layers represent different relationships that sub-word tokens can have with each other. As it learns, the model updates its understanding of both of these relationships at the same time. And during inference, it will consider what it has learned about both of these relationships when it predicts the next most likely token to occur.

Fine-tuning Process

The fine-tuning process involves adjusting the pre-trained DistilBERT model to our specific classification task.

The code below is implemented using Hugging Face’s Transformers library, which provides a high-level interface for working with DistilBERT and other transformer models. This interface includes pre-built model architectures, tokenizers, and data processing utilities. The library offers components like DistilBertTokenizer for text tokenization and DistilBertForSequenceClassification for the actual model architecture. It also provides Dataset and DataLoader classes for efficient data handling and batching.

These high-level components abstract away much of the complexity involved in working with transformer models. The tokenizer handles the task of converting raw text into a format the model can understand, including subword tokenization and special token management. The model class encapsulates DistilBERT’s architecture, including the self-attention mechanisms and feed-forward networks, exposing simple methods for forward passes and loss computation.

Transformers delegates the numerical computations to PyTorch, a deep learning framework that handles low-level operations. PyTorch manages tensor operations, automatic differentiation for backpropagation, and GPU acceleration. It provides the foundation for defining and training neural networks, offering classes like torch.nn for neural network layers and torch.optim for optimization algorithms.

Data Preparation

The process begins by loading the labeled data. It consisted of three files, one for each category, each containing several hundred customer messages.

def load_data(file_path):
    with open(file_path, 'r') as f:
        return [line.strip() for line in f]

feedback = load_data('data/feedback.txt')
inquiries = load_data('data/inquiries.txt')
refunds = load_data('data/refunds.txt')

all_texts = feedback + inquiries + refunds
all_labels = [0] * len(feedback) + [1] * len(inquiries) + [2] * len(refunds)

This code loads the text strings and assigns numerical labels to each. In neural networks, labels refer to “the right answer” and are used during training to check the model’s guesses.

Dataset and DataLoader

A custom dataset class is created to handle tokenization:

class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

    def __getitem__(self, item):
        text = str(self.texts[item])
        label = self.labels[item]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )

        return {
            'text': text,
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

A tokenizer is a tool that breaks down text into smaller units called tokens. These could be words, parts of words, or even punctuation. The tokenizer also handles converting these tokens into numbers that the model can understand. The DistilBertTokenizer is used because the same tokenizer that the model was trained with is desired. Under the hood, it uses WordPiece subword segmentation.

The encode_plus method of this tokenizer is used. This method does several things: It tokenizes the input text, adds special tokens that DistilBERT expects (like [CLS] at the start and [SEP] at the end), pads or truncates the input to a specified maximum length, and creates an “attention mask” which tells the model which tokens are actual input and which are padding.

The torch Dataset is a PyTorch class that represents a dataset. By inheriting from this class, a custom dataset that PyTorch can easily work with is created. The __getitem__ method is called when an item from the dataset needs to be accessed.

PyTorch’s DataLoader is used to efficiently batch and shuffle the data:

train_dataloader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=16)

The DataLoader is a PyTorch utility that helps manage batching of data and provides an iterable over the Dataset. For the training data, shuffle=True is set, which means it will randomly shuffle the data at the start of each epoch. This shuffling helps prevent the model from learning any unintended patterns based on the order of the training data.

A batch size of 16 is used, meaning the DataLoader will yield batches of 16 examples at a time. This batch size is a balance between memory usage and training speed. For the test data, shuffling is not needed, so the shuffle parameter is omitted.

During training and evaluation, these DataLoaders will be iterated over, which will give batches of data in the format the model expects. This abstraction simplifies the training loop and makes it easier to work with large datasets that might not fit into memory all at once.

Train-Test Split

Before training begins, the data is split into training and testing sets:

def train_test_split(data, test_size=0.1):
    split_index = int(len(data) * (1 - test_size))
    return data[:split_index], data[split_index:]

train_data, test_data = train_test_split(list(zip(all_texts, all_labels)))
train_texts, train_labels = zip(*train_data)
test_texts, test_labels = zip(*test_data)

A manual split is used where the first 90% of the data is used for training and the last 10% for testing. This approach was chosen to make it easy for me to place all of the samples evaluated in the previous blog post into the test set. That way, performance could be more directly compared, because none of those samples would have been seen by any either model during training. More on that later.

Training Loop

The training process involves iterating over the data multiple times and updating the model’s parameters to minimize the classification error.

def train_epoch(model, data_loader, optimizer, device):
    model.train()
    total_loss = 0
    total_correct = 0
    total_samples = 0

    for batch in data_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        logits = outputs.logits

        _, predicted = torch.max(logits, 1)
        total_correct += (predicted == labels).sum().item()
        total_samples += labels.size(0)

        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    avg_loss = total_loss / len(data_loader)
    accuracy = total_correct / total_samples
    return avg_loss, accuracy

It isn’t shown here, but model passed in to this function is an instance of the DistilBertForSequenceClassification class. input_ids, attention mask, and labels are determined during dataset preparation outside of this function call. The input_ids represent the tokenized input text, while attention_mask indicates which tokens should be attended to. labels represent the correct answers, while predicted represents the guesses.

The function performs optimization - the process where the model adjusts its parameters to minimize the loss and iteratively improve its performance on the task. This involves a “forward pass” where the model processes the input and makes predictions. The loss is then calculated by comparing these predictions to the correct labels. Next, a “backward pass” (backpropagation) computes the gradients. Finally, the optimizer uses these gradients to update the model’s parameters, adjusting weights to reduce the loss.

Calling model in this manner is the PyTorch idiom for doing a forward pass (this has always been a strange design choice to me).

The backward pass is done by a call to loss.backward() where it computes the direction and amount to nudge its parameters and optimizer.step() where it makes those parameter updates.

Most of the actual number crunching is carried out by PyTorch’s autograd code. The attention mechanics are implemented in DistilBertForSequenceClassification class of the Hugging Face transformers library. This library also provides AdamW’s ability to perform parameter updates (with additional optimizations, in AdamW’s case).

As you can see, there isn’t really a defined API for gradient descent. It’s on you to know what steps need to be done and do them.

Model Configuration and Training

The model configuration and training takes place in the program’s main method, separate from the above helper methods.

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=3)

train_dataset = TextClassificationDataset(train_texts, train_labels, tokenizer, max_len=128)
test_dataset = TextClassificationDataset(test_texts, test_labels, tokenizer, max_len=128)

First, the tokenizer and model are initialized. The pre-trained ‘distilbert-base-uncased’ model is used, which means it’s been trained on lowercase English text. The num_labels=3 parameter tells the model it’s dealing with a three-class classification problem.

max_len=128 is set when creating the datasets. This maximum length parameter determines the longest sequence of tokens the model will process. Shorter sequences will be padded, and longer ones will be truncated.

The batch size of 16 in the DataLoader means 16 examples will be processed at a time during training. This is another balance between memory usage and training speed. Larger batch sizes can lead to faster training but require more memory.

device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
model.to(device)

Next, the model is moved to the appropriate device. A check is made if Apple’s Metal Performance Shaders (MPS) are available, which can accelerate training on compatible Mac hardware (this was only run on a Mac; this line would need to be changed if it were to be run optimally on different hardware).

optimizer = AdamW(model.parameters(), lr=2e-5)

AdamW is an optimization algorithm for training neural networks. It features adaptive learning rates for each parameter and implements decoupled weight decay regularization. Weight decay is a regularization technique that encourages smaller parameter values in the model. It does this by adding a penalty term to the loss function based on the magnitude of the weights. The primary goal of weight decay is to prevent overfitting by discouraging the model from relying too heavily on any individual feature or learning overly complex patterns that may not generalize well to new data.

A learning rate of 2e-5 (0.00002) is used here. It was selected by manually tuning the learning rate and number of epochs and evaluating network performance. It’s small enough to allow for fine adjustments to the pre-trained weights without causing drastic changes that could destroy the model’s pre-trained knowledge. Too high, and the model might overshoot optimal solutions; too low, and the model might train too slowly or get stuck in suboptimal solutions. This learning rate is much lower than what you would use for training a network from scratch, which is usually around 1e-3.

num_epochs = 3
for epoch in range(num_epochs):
    print(f'Epoch {epoch + 1}/{num_epochs}')
    train_loss, train_acc = train_epoch(model, train_dataloader, optimizer, device)
    print(f'Train loss {train_loss:.4f} accuracy {train_acc:.4f}')

    val_loss, val_acc = evaluate(model, test_dataloader, device)
    print(f'Val loss {val_loss:.4f} accuracy {val_acc:.4f}')

torch.save(model.state_dict(), 'model/distilbert_classifier.pth')
print("Model saved successfully.")

The training loop runs for 3 epochs, after which the accuracy and loss of both the training and validation sets reached desirable values (see “Results and Discussion” below).

An epoch is one complete pass through the entire training dataset. In each epoch, the model is trained on the training data and then evaluated on the validation (test) data. The loss and accuracy for both training and validation sets are printed out with 4 decimal places.

After training, the model is saved using torch.save(). This function saves the model’s state_dict. A state_dict in PyTorch is a Python dictionary that maps each layer to its parameter tensors. It contains all the learned weights and biases of the model. By saving the state_dict, all the knowledge the model has gained during training is essentially saved. The model is saved with a .pth file extension. This is a common convention in PyTorch for saved model files, standing for “PyTorch”.

To use this saved model later, it would be loaded like this:

model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased', num_labels=3)
model.load_state_dict(torch.load('model/distilbert_classifier.pth'))
model.eval()

Results and Discussion

For simplicity, I ran with the assumption that each string of text had one and only one classification. It would be much more useful if each string of text instead could have 0-3 classifications. This would also almost certainly bring down the accuracy of the model because determining whether a given string belongs to a single category at all seems like it would be a much more complex task than what it had to learn during training.

After 3 epochs, the model achieves an accuracy of 100% and loss of 0.156 on the test set and an accuracy of 99.59% and loss of 0.0278 on the training set.

Epoch 1/3
Train loss 0.6052 accuracy 0.8264
Val loss 0.1448 accuracy 1.0000
Epoch 2/3
Train loss 0.0708 accuracy 0.9959
Val loss 0.0314 accuracy 1.0000
Epoch 3/3
Train loss 0.0278 accuracy 0.9959
Val loss 0.0156 accuracy 1.0000

The task was designed to be easy enough that a network could handle it, so the performance met my expectations. The 100% accuracy of the test set includes the entirety of the tasks that Llama 3 8B only achieved 92% accuracy on, with additional test data added in. And just for kicks, I later tried coming up with weird and contrived strings of text to try to confuse the model. I could not get it to incorrectly classify anything.

After the fine-tuned DistilBERT model was trained on learning the distinctions between the three specific categories (refund requests, order inquiries, and general feedback), the fine-tuned model outperformed base Llama 3 for this specialized task at roughly 0.8% of its size. Llama 3 8B is subject to the variability of a general purpose instruction model, so it makes sense that it gets tripped up on even simple tasks from time to time.

I came away completely sold on the idea of fine-tuning small models for repetitive, special purpose tasks. It is inevitable that many suitable tasks will instead be run on models that are hundreds of billions of parameters or larger via commercial APIs. This presents a huge opportunity for cost savings and efficiency that will only grow over time.

The source code for this project is available on GitHub.

Update 2025-05-18: Looking back, this post is a little embarassing, but I’ll leave it up in case any other learners stumble across it and find it helpful.

In the video The spelled-out intro to language modeling: building makemore, Andrej Karpathy briefly mentioned a bug that can occur when broadcasting tensors during mathematical operations.

In the lines below, N is a 27x27 tensor. The code is trying to normalize the rows of the tensor by dividing each row by the sum of the row.

Here is the incorrect calculation:

P = N.float()
P /= P.sum(1)

The correct calculation should be:

P = N.float()
P /= P.sum(1,keepdim=True)

Andrej briefly explained why this happens, but his explanation didn’t land for me. I walked through a few smaller examples to understand the mechanics of the tensor operations.

P / P.sum(1, keepdim=True)

P = [[1, 2, 3],
     [4, 5, 6],
     [7, 8, 9]]

P = torch.tensor(P)
P_sum = torch.tensor(P).sum(1, keepdim=True)

# P_sum: this is what a column vector looks like in PyTorch
# [[6],
#  [15],
#  [24]]

Here, P.sum(1, keepdim=True) computes the sum of each row in P and returns its results as a 3 element column vector. The 1 argument is specifying that we want sums computed by the first axis, which are the rows of the vector. keepdim=True is saying that the output should have the same dimensions as the input. In other words, each row should remain a vector, rather than being flattened during the summation operation.

Now when you divide P by P_sum, each element in a row of P is divided by the sum of that row. This normalizes each row so that the sum of the elements in each row is 1.

P / P.sum(1, keepdim=True)  

# [[1/6, 2/6, 3/6],
# [4/15, 5/15, 6/15],
# [7/24, 8/24, 9/24]]

P / P.sum(1)

P = [[1, 2, 3],
     [4, 5, 6],
     [7, 8, 9]]
     
P = torch.tensor(P)
P_sum = P.sum(1) 

# [6, 15, 24]

P.sum(1) computes the sum of each row in P but it flattens the results into a single, flat row vector. The inner dimensions are discarded.

When you divide P by P.sum(1), the broadcasting rules in PyTorch will try to align the shapes. The broadcasting will align the 3 elements of P.sum(1) with the columns of P, which is not the intended behavior. This will result in incorrect normalization. Notice how each summed row is dividing a row, rather than a column. It’s operating on the wrong dimension.

P / P.sum(1)  

# [[1/6, 2/15, 3/24],
#  [4/6, 5/15, 6/24],
#  [7/6, 8/15, 9/24]]

PyTorch mechanics can seem a little unintuitive and mind-bendy when you’re learning them, so it can be helpful to pick them apart like we did here to understand them.

I’ve encountered many tasks that could almost be automated, except that they have some kind of ambigious step that requires manual attention. For example, it’s extremely common to find instances where freeform text needs to be flagged if it contains sensitive or important information. For such tasks, it’s easy to write software that performs subsequent steps such as interacting with a database or moving files around, but ambiguous text-processing can be extremely difficult to implement programatically. Solutions that use regular expressions or other simple text processing techniques can be finnicky, inaccurate, and difficult to tune.

LLMs have presented an opportunity to make big gains in this problem space. Most programming languages can easily call into large commercial models like OpenAI’s GPT-4 or Google’s Gemini. You can read text in from your program, pass it to these service’s API, and perform actions based on the results. However, using external services for this can be problematic because of the cost and privacy implications. You may not want to depend on a paid service for an important task, and you really may not want to pass organizational data to it.

On the other hand, open source LLMs that run locally on your system have been rapidly improving. They aren’t quite as powerful as commercial models, but they can be surprisingly effective, particularly if you give them simple tasks and make efforts to constrain their responses. Additionally, you always have the option of fine-tuning to improve performance.

One of the biggest open source project to run these models is llama.cpp, which performs inference on open source models. It can be built as a library and utilized from other languages which write bindings to its API. I really like to experiment with new problem spaces in Clojure, so I used llama.clj to experiment with programatically interfacing with LLMs to automate tasks.

I decided to put together an example to experiment with to simulate a case where we have a bunch of customer service inquiries and we want to figure out whether they are refund requests, order inquiries, or just general feedback. The goal of these experiments was only to perform this classification step.

(require '[com.phronemophobic.llama :as llama])
(require '[com.phronemophobic.llama.util :as llutil])

;; 8B parameter llama 3 model with 4bit quantization that easily runs on my Macbook M1
(def model-path "/Users/jon/development/cpp/llama.cpp/models/Meta-Llama-3-8B-Instruct.Q4_0.gguf")
(def llama-context (llama/create-context model-path {}))

(def inquiries
  [{:classification "Order Inquiry" :inquiry "Where is my order? It was supposed to arrive yesterday."}
   {:classification "Refund Request" :inquiry "I want to return an item and get a refund. Can you help me with that?"}
   {:classification "General Feedback" :inquiry "I think your website could be more user-friendly."}
   ;; ... many more omitted
])

(def classifications ["Order Inquiry" "Refund Request" "General Feedback"])

My setup was to run each classification over a classify function and append the experimental result to the actual result so that I could compare them.

(defn naive-classify [inquiry]
  (llama/generate-string
   llama-context
   (llama3-inquiry inquiry)))

(defn naive-greedy-classify [inquiry]
  (llama/generate-string
   llama-context
   (llama3-inquiry inquiry)
   {:samplef llama/sample-logits-greedy}))

(defn classify-inquiries [classify-fn inquiries]
  (map (fn [inquiry]
         (assoc inquiry
                :experimental-classification
                (classify-fn (:inquiry inquiry))))
       inquiries))
       
(defn correct-results [results]
  (filter #(= (:classification %) (:experimental-classification %)) results))

(defn incorrect-results [results]
  (filter #(not= (:classification %) (:experimental-classification %)) results))

(defn summary-results [results]
  {:correct (count (correct-results results))
   :incorrect (count (incorrect-results results))
   :accuracy (/ (count (correct-results results)) (count results))})

llama.clj provides a few different sampling functions that you can use to generate text. Sampling functions are the means by which the text tokens which form the response are selected. In addition to giving a few options for pre-defined functions, the library also gives you the ability to create your own. For my experiments, I started with microstatv2 (which naive-classify above uses above - llama.clj uses this function by default) and a greedy sampling function (which naive-greedy-classify requests via the samplef key).

I noticed I would generally get about 92% accuracy with the model, data, and prompts I used. When I looked at the responses that were incorrect, I noticed that many of them were from responses that would misformat the specific response format that I asked for, e.g. returning something like “[Order Inquiry]” instead of just “Order Inquiry”. Since llama.cpp is feeding you text straight from the model, you have to get a little creative to deal with this. Both llama.cpp and llama.clj provide different ways to constrain models to produce valid json, which you can then further validate. I decided to try constraining responses by first just retrying if the gets a response does not belong to the set of strings I’m looking for:

(defn retry-classify [inquiry]
  (loop [count 0]
    (let [response (llama/generate-string
                    llama-context
                    (llama3-inquiry inquiry))]
      (if (some #(= response %) classifications)
        response
        (if (> count 3)
          nil
          (recur (inc count)))))))

This took care of those cases nicely and brought accuracy up to about 95%. It didn’t noticeably decrease performance because it doesn’t retry often - only in specific failure cases.

One potential issue is that if the model ever went haywire and started filling the context window with garbage (which does happen occasionally on some models), it would take inordinate amounts of time to generate responses and fail, which would slow everything to a crawl. I decided to try hand rolling a greedy sampling function that selects the first token which completes a valid classification to avoid this situation. The way that llama.cpp works is that it determines the relative probability of every possible text token at each point in the response. What I wanted to try to do in my sampling function was constrain the responses by forcing it to only choose tokens that are available in pre-defined categories (defined by classifications above). For the first generated token, it can only select the first token of any of those classifications. It selects the member of those tokens which it deems most likely to be correct.

From there, at each step, we filter out only classifications which start with the response tokens we’ve already accumulated, and repeat the process until we have a full response.

(defn greedy-constrained-classify [inquiry]
  (let [prompt (llama3-prompt (str "Inquiries can be one of 'Order Inquiry', 'Refund Request', or 'General Feedback'. What is the classification of the following inquiry? Reply with only the classification and nothing else: \"" inquiry "\""))
        prompt-tokens (llutil/tokenize llama-context prompt)]
    (llama/llama-update llama-context (llama/bos) 0)
    (doseq [token prompt-tokens]
      (llama/llama-update llama-context token))
    (loop [classification-tokens (map #(llutil/tokenize llama-context %) classifications)
           acc []]
      (let [logits (llama/get-logits llama-context)
            valid-tokens (map first classification-tokens)
            token (->> logits
                       (map-indexed (fn [idx p]
                                      [idx p]))
                       (filter #(contains? (set valid-tokens) (first %)))
                       (apply max-key second)
                       first)
            next-acc (conj acc token)
            next-classification-tokens (->> classification-tokens
                                            (filter #(= (first %) token))
                                            (map rest)
                                            (map vec)
                                            (into []))]
        (if (or (empty? next-classification-tokens) (every? empty? next-classification-tokens))
          (llutil/untokenize llama-context next-acc)
          (do
            (llama/llama-update llama-context token)
            (recur next-classification-tokens next-acc)))))))

Unfortunately, in my testing, this brought accuracy way down to around 84%. When I dug into the issue, it seemed that in failing cases, the very first response token it generates is for the wrong classification. It’s kind of hard for me to tell what the issue is without more knowledge about llama.cpp or how the wrapper works. It’s possible that that there may be an issue with the initial context setup while reading in the prompt. I emulated the setup steps that I found in llama.clj’s documentation, but perhaps llama.cpp has changed since the repository was written (the examples and tutorials were written with llama 2), or maybe llama 3 is fundamentally different and requires different processing from llama 2, or maybe I overlooked some other mistake in my code.

I’m planning to dig into llama.cpp more in depth next to learn more and continue experimenting. By looking at llama.cpp’s internal mechanics, it would be interesting to see why my sampling function failed. Before that, I think it might be a good idea to work through Andrej Karpathy’s neural network course to try to get some more context and understanding of how these libraries work under the hood. It probably isn’t strictly necessary, but the more domain knowledge you have, the easier it is to understand the architecture of software projects.

Overall, llama.clj is extremely fun to learn and sketch out ideas with. Even without much knowledge about LLMs, you could probably use it put together some interesting and effective applications. I came in expecting just a dry but functional wrapper to llama.cpp, but the documentation was fantastic and I learned a lot of new techniques and concepts from the clojure code used throughout the project. From the little bit of llama.cpp that I looked at, it didn’t really seem to strive to maintain API stability, so it seems like it must be a challenge for wrapper libraries such as llama.clj to keep up to date. In any case, I had a great time and I’m looking forward to learning more about training and inference.

ollama is a software project which makes it easy to run LLMs on your local machine. Running ollama run llama3 downloaded a 4-bit quantized model that could run on my Macbook M2, and then dropped me into a CLI where I could enter prompts and responses stream into my terminal.

In the above clip, I’m running Meta’s latest open source model Llama 3. These models are less powerful than OpenAI’s models - which are mainly GPT-4o, GPT-4, and GPT-3.5 Turbo as of writing, but they pack a serious punch. The fact that a model that can run on a laptop can get within throwing distance of GPT-3.5, which powered all of ChatGPT not long ago is pretty insane.

There’s a number of reasons you would want to run your own local AI instead of something like ChatGPT:

1. Privacy. With ChatGPT, OpenAI has access to every query and response you send. When someone else controls your data, despite their best intentions, sometimes it will leak. In OpenAI’s case, it already has already happened at least once. Depending on the sensitivity of your prompts, this may be more or less of a concern.
2. Flexibility. Using your own models allows you to use special purpose models better suited for individual tasks. You can also fine-tune open source models on your own hardware, which might be incredibly useful if you have a lot of organizational data you would like the model to recognize.

The project’s README shows some of its capabilities:

REST API

Ollama has a REST API for running and managing models.

Generate a response

curl http://localhost:11434/api/generate -d '{
  "model": "llama3",
  "prompt":"Why is the sky blue?"
}'

Chat with a model

curl http://localhost:11434/api/chat -d '{
  "model": "llama3",
  "messages": [
    { "role": "user", "content": "why is the sky blue?" }
  ]
}'

See the API documentation for all endpoints.

We can try one of those curl commands and see what the responses look like:

{"model":"llama3","created_at":"2024-06-20T00:25:35.629748Z","response":"The","done":false}
{"model":"llama3","created_at":"2024-06-20T00:25:35.662856Z","response":" sky","done":false}
{"model":"llama3","created_at":"2024-06-20T00:25:35.695868Z","response":" appears","done":false}
{"model":"llama3","created_at":"2024-06-20T00:25:35.729704Z","response":" blue","done":false}
{"model":"llama3","created_at":"2024-06-20T00:25:35.763045Z","response":" because","done":false}
...

These stream in over time instead of being dumped out when the message is complete. Here’s the code from ollama that’s generating the stream:

func streamResponse(c *gin.Context, ch chan any) {
	c.Header("Content-Type", "application/x-ndjson")
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
			slog.Info(fmt.Sprintf("streamResponse: json.Marshal failed with %s", err))
			return false
		}

		// Delineate chunks with new-line delimiter
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
			slog.Info(fmt.Sprintf("streamResponse: w.Write failed with %s", err))
			return false
		}

		return true
	})
}

application/x-ndjson means “Newline delimited JSON”. This detail makes it easier to delineate between messages. Since you know they are broken on newlines, you know you can use line-reading functionality that’s present in many languages to process them.

A Simple ChatGPT Clone

For fun, we can build a simple ChatGPT-style web application on top of this API with a few components:

• an HTML page that presents a form to collect a prompt with a button for submission, and javascript to send the prompt, listens for a response, and writes the responses to the page as they stream in
• a server with two endpoints:
  • one rendering the page above,
  • one handling receiving a prompt, marshalling it to ollama, and streaming the response to the client.

With that in mind, let’s get to work!

Go Implementation

We can do everything in Go using just the standard library.

package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"log"
	"net/http"
)

func main() {
	tmpl := template.Must(template.ParseFiles("index.html"))

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if err := tmpl.Execute(w, nil); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	})

	http.HandleFunc("/api/generate", func(w http.ResponseWriter, r *http.Request) {
		var request struct {
			Model  string `json:"model"`
			Prompt string `json:"prompt"`
		}
		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		payload := map[string]string{
			"model":  request.Model,
			"prompt": request.Prompt,
		}

		payloadBytes, err := json.Marshal(payload)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		resp, err := http.Post("http://localhost:11434/api/generate", "application/json", bytes.NewBuffer(payloadBytes))
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer resp.Body.Close()

		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)

		decoder := json.NewDecoder(resp.Body)
		for {
			var response map[string]interface{}
			if err := decoder.Decode(&response); err == io.EOF {
				break
			} else if err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			if err := json.NewEncoder(w).Encode(response); err != nil {
				http.Error(w, err.Error(), http.StatusInternalServerError)
				return
			}
			w.(http.Flusher).Flush()
		}
	})

	fmt.Println("Server is running on http://localhost:8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}

Handling the chunked JSON is done in that last block with the decoder. The decoder is a state machine that handles its own buffering. The for loop checks for error conditions, and if they aren’t found, will flush the buffer contents to the client, which in this case is a javascript program running inside the browser. As the Javascript program receives those messages, it writes them on the page.

Decoding and immediately re-encoding the result looks a little silly, we could have just pointed the client directly to ollama and left all of this out of the server. But in a more realistic deployment, we would not want clients interacting directly with the model - it would be important for the server to sit in between and manage the process.

That was all it took to build the Go implementation! I like how skimming the docs for standard library modules that sounded like what you’re looking for is all it really takes to get up and running in the language. The fact that I can do that when I’m not very experienced in the language is a testament to the skill of its design.

The rest of this post will cover building this same functionality in Clojure instead of Go.

Clojure - HTTP kit Implementation

Unlike Go, Clojure does not have a production grade HTTP server and client available in the standard library.

I used http-kit for this task because it was the smallest library I knew of which could single handedly meet my requirements - it has an synchronous/asynchronous http clients as well as synchronous/asynchronous Ring-compliant web servers. I also chose it because of its reputation of being small and focused with minimal dependencies.

In Clojure, Ring-compliant HTTP server’s use handlers that are maps containing keys for various http constructs such as :status, :headers, and :body. I used this as a starting point and built a run of the mill synchronous handler that serves the same page as the Go app above.

(defn read-html-template []
  (println "fetching html template")
  (slurp (io/resource "index.html")))

(defn index-handler [req]
  (println "in index handler")
  {:status  200
   :headers {"Content-Type" "text/html"}
   :body    (read-html-template)})
   
(defn not-found-handler [req]
  {:status  404
   :headers {"Content-Type" "text/plain"}
   :body    "Page not found."})
   
(defn app [req]
  (let [uri (:uri req)
        method (:request-method req)]
    (cond
      (and (= uri "/") (= method :get)) (index-handler req)
      :else (not-found-handler req)))) 

That all works because you can pass any arbitrary string as your response body and set its content type to text/html, it will render correctly.

At this point, the client has a javascript programming that wants to pass a prompt and listen for the streaming response. We will want an asynchronous handler to do this, so we’ll use http-kit’s as-channel function. It’s not covered in the project’s main documentation, but the docstring in the source itself has enough to get started with. Let’s use it and wire it up.

(def clients_ (atom #{}))
(defn my-async-handler [ring-req]
  (http/as-channel ring-req
              {:on-open (fn [ch]
                          (println "conn open!")
                          (println ring-req)
                          (swap! clients_ conj ch))
               :on-close (fn [ch]
                           (println "conn close!")
                           (swap! clients_ disj ch))}))

(defn app [req]
  (let [uri (:uri req)
        method (:request-method req)]
    (cond
      (and (= uri "/") (= method :get)) (index-handler req)
      (and (= uri "/api/generate") (= method :post)) (my-async-handler req)
      :else (not-found-handler req))))

When I evaluate these forms, switch back to my browser, and hit the “Submit” buttton, I see the print statements fire. When I evaliate clients_, I see it now has a new client.

To get the actual messages, we can start writing a function to send a POST request to the ollama server. After that, we want to find a way to access messages as they stream in and send them to the client.

In typical asynchronous programming, this kind of message passage is done with callbacks. http-kit has its own concept of channels that is separate from core.async with its own semantics. When we get a message, we will want to send it to the client like this:

(http/send! ch {:status 200
                :headers {"Content-Type" "application/json"}
                :body json-encoded}
            false)

To kick off the HTTP request to ollama with the prompt, we want to use http-kit’s client functionality. There is documentation on making asynchronous requests with callbacks:

;fire and forget, returns immediately[1], returned promise is ignored
(http/get "http://host.com/path")

(def options {:timeout 200             ; ms
              :basic-auth ["user" "pass"]
              :query-params {:param "value" :param2 ["value1" "value2"]}
              :user-agent "User-Agent-string"
              :headers {"X-Header" "Value"}})
(http/get "http://host.com/path" options
          (fn [{:keys [status headers body error]}] ;; asynchronous response handling
            (if error
              (println "Failed, exception is " error)
              (println "Async HTTP GET: " status))))
 ; [1] may not always true, since DNS lookup maybe slow

This will kind of work, but it will fire the promise once all the messages are received. We’re looking for some way to access the messages as the stream in.

I looked at the source for help and found the following. Without guidance from the documentation, the :stream option sounded like a good option to try.

   Returned body type is controlled by `:as` option:

    Without automatic unzipping:
      `:none`           - org.httpkit.DynamicBytes
      `:raw-byte-array` - bytes[]

    With automatic unzipping:
      `:byte-array`     - bytes[]
      `:stream`         - ByteInputStream
      `:text`           - String (charset based on Content-Type header)
      `:auto`           - As `:text` or `:stream` (based on Content-Type header)

There’s no further mention of how to use a ByteInputStream in the docs, so we can check the source for that.

package org.httpkit;

import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;

/**
 * No synchronization, better toString
 */
public class BytesInputStream extends InputStream {
    private final byte[] buf;
    private final int count;
    private int mark = 0;

    private int pos;

    public BytesInputStream(byte[] data, int length) {
        this.buf = data;
        this.count = length;
        this.pos = 0;
    }

    /**
     * get the underlying bytes, copied
     *
     * @return
     */
    public byte[] bytes() {
        return Arrays.copyOf(buf, count);
    }

    public int read() throws IOException {
        return (pos < count) ? (buf[pos++] & 0xff) : -1;
    }
...

Since this class subclasses InputStream, it can be read by an InputStreamReader.

The docs for InputStreamReader recommend wrapping with a BufferedReader, so let’s do that. We want to take inputs line by line

(defn send-prompt-to-ollama [ch prompt]
  (let [url "http://localhost:11434/api/generate"]
    (client/post url
                 {:as :stream
                  :headers {"Content-Type" "application/json"}
                  :body (json/write-str {:model "llama3" :prompt prompt})}
                 (fn [{:keys [status headers body error]}]
                   (if error
                     (do
                       (http/send! ch {:status 500 :body (str "Internal Server Error: " error)})
                       (http/close ch))
                     (let [stream ^java.io.InputStream body
                           reader (java.io.BufferedReader. (java.io.InputStreamReader. stream "UTF-8"))]
                       (loop []
                         (let [line (.readLine reader)]
                           (if (nil? line)
                             (do
                               (.close reader)
                               (http/close ch))
                             (do
                               (let [response (:response (json/read-json line))
                                     json-encoded (json/write-str {:response response})]
                                 (http/send! ch {:status 200
                                                 :headers {"Content-Type" "application/json"}
                                                 :body json-encoded}
                                             false))
                               (recur)))))))))))

Unfortunately, after writing this up, I observed the same behavior as earlier: responses were not streamed in, but rather returned all at once.

It turns out this appears to be a known limitation of http-kit’s client functionality. The comment thread mentioned that clj-http works for this use case, but part of the reason I went with http-kit in the first place was to minimize dependencies.

Another option to try without introducing a new dependency is to interop with java.net.http.HttpClient, which has been included with the JDK since Java 11.

Here’s how I initially got that working:

(import (java.net.http HttpClient HttpRequest HttpResponse HttpResponse$BodyHandlers HttpRequest$BodyPublishers)
        (java.net URI)
        (java.nio.charset StandardCharsets)
        (java.io InputStreamReader BufferedReader)
        (java.util.concurrent CompletableFuture))
        
(defn handle-response [ch response]
  (with-open [reader (BufferedReader. (InputStreamReader. (.body response) StandardCharsets/UTF_8))]
    (loop []
      (let [line (.readLine reader)]
        (if (nil? line)
          (do
            (.close reader)
            (http/close ch))
          (do
            (println "got line" line)
            (let [response (:response (json/read-json line))
                  json-encoded (json/write-str {:response response})]
              (println "sending response" response)
              (http/send! ch {:status 200
                              :headers {"Content-Type" "application/json"}
                              :body json-encoded}
                          false))
            (recur)))))))

(defn send-async-request [ch model prompt]
  (let [client (HttpClient/newHttpClient)
        body (str "{\"model\":\"" model "\", \"prompt\":\"" prompt "\"}")
        request (-> (HttpRequest/newBuilder)
                    (.uri (URI/create "http://localhost:11434/api/generate"))
                    (.header "Content-Type" "application/json")
                    (.POST (HttpRequest$BodyPublishers/ofString body))
                     (.build))]
    (-> (.sendAsync client request (HttpResponse$BodyHandlers/ofInputStream))
        (.thenAccept (reify java.util.function.Consumer
                       (accept [_ response]
                         (handle-response ch response))))
        (.exceptionally (reify java.util.function.Function
                          (apply [_ error]
                            (println (str "Request failed: " error))
                            nil))))))

This gave me the behavior I was looking for: responses were sent to the client as soon as they come in.

The code is a little noisy because the library uses Java idioms that were introduced after Clojure was created.

Luckily, Clojure’s maintainers recently introduced language changes that make using these features less painful. From the release notes of Clojure 1.12.0-alpha12:

Functional interfaces

Java programs define "functions" with Java functional interfaces (marked with the @FunctionalInterface annotation), which have a single method.

Clojure developers can now invoke Java methods taking functional interfaces by passing functions with matching arity. The Clojure compiler implicitly converts Clojure functions to the required functional interface by constructing a lambda adapter. You can explicitly coerce a Clojure function to a functional interface by hinting the binding name in a let binding, e.g. to avoid repeated adapter construction in a loop.

See: CLJ-2799

To take advantage of these changes, we can use a preview build of the language by updating the language version in deps.edn:

org.clojure/clojure       {:mvn/version "1.12.0-beta1"}

After doing so, we can change the above functions to this:

(defn handle-response [ch response]
  (with-open [reader (BufferedReader. (InputStreamReader. (.body response) StandardCharsets/UTF_8))]
    (loop []
      (let [line (.readLine reader)]
        (if (nil? line)
          (do
            (.close reader)
            (http/close ch))
          (let [response (:response (json/read-json line))
                json-encoded (json/write-str {:response response})]
            (http/send! ch {:status 200
                            :headers {"Content-Type" "application/json"}
                            :body json-encoded}
                        false))
          (recur))))))

(defn send-async-request [ch model prompt]
  (let [client (HttpClient/newHttpClient)
        body (str "{\"model\":\"" model "\", \"prompt\":\"" prompt "\"}")
        request (-> (HttpRequest/newBuilder)
                    (.uri (URI/create "http://localhost:11434/api/generate"))
                    (.header "Content-Type" "application/json")
                    (.POST (HttpRequest$BodyPublishers/ofString body))
                    (.build))]
    (-> (.sendAsync client request (HttpResponse$BodyHandlers/ofInputStream))
        (.thenAccept (fn [response] (handle-response ch response)))
        (.exceptionally (fn [error] (println (str "Request failed: " error)))))))

From here, I just needed to wire up the request logic to my async handler and then update my main handler to route to it.

(defn my-async-handler [ring-req]
  (let [body (slurp (:body ring-req))
        prompt (try
                 (:prompt (json/read-json body))
                 (catch Exception e
                   (println "Error parsing request body:" e)
                   nil))]
    (if prompt
      (http/as-channel ring-req
                       {:on-open (fn [ch]
                                   (println "conn open!")
                                   (swap! clients_ conj ch)
                                   (send-async-request ch "llama3" prompt))
                        :on-close (fn [ch]
                                    (println "conn close!")
                                    (swap! clients_ disj ch))})
      {:status 400
       :headers {"Content-Type" "application/json"}
       :body (json/write-str {:error "Invalid request"})})))

(defn app [req]
  (let [uri (:uri req)
        method (:request-method req)]
    (cond
      (and (= uri "/") (= method :get)) (index-handler req)
      (and (= uri "/api/generate") (= method :post)) (#'my-async-handler req)
      :else (not-found-handler req))))

Clojure - Pedestal Implementation

I was hoping this project would be a good place to use core.async, which I just started recently using, but realized while working through the previous implementation that http-kit has its own semantics for channels. It didn’t really make sense to use core.async there.

I heard that Pedestal was built with core.async in mind, so I ended up throwing together another implementation with it.

In this version, my request function sends a request with the user’s prompt and returns a channel that has the messages queued.

(defn handle-response [response-ch ^HttpResponse response]
  (let [reader (BufferedReader. (InputStreamReader. (.body response) StandardCharsets/UTF_8))]
    (go-loop []
      (if-let [line (.readLine reader)]
        (do
          (println line)
          (a/>! response-ch line)
          (recur))
        (a/close! response-ch)))))

(defn send-async-request [model prompt]
  (let [client (HttpClient/newHttpClient)
        body (generate-string {:model model :prompt prompt})
        request (-> (HttpRequest/newBuilder)
                    (.uri (URI/create "http://localhost:11434/api/generate"))
                    (.header "Content-Type" "application/json")
                    (.POST (HttpRequest$BodyPublishers/ofString body))
                    (.build))
        response-ch (a/chan)]
    (-> (.sendAsync client request (HttpResponse$BodyHandlers/ofInputStream))
        (.thenAccept (fn [response] (handle-response response-ch response)))
        (.exceptionally (fn [error] (println (str "Request failed: " error)))))
    response-ch))

I used Pedestal’s Server Sent Events functionality to make it work. The json interceptor is necessary to get access to the json params sent from the client to the asynchronous endpoint.

start-event-stream returns an interceptor and a function to call which takes a channel as an argument - you’ll put messages on it to send them to the client. The channel’s buffer is maintained by the library. In my callback function, I pop messages off the request channel and put them onto the request channel.

(defn stream-ready [event-ch ctx]
  (let [{:keys [model prompt] :as raw} (-> ctx :request :json-params)
        response-chan (send-async-request model prompt)]
    (go-loop []
      (if-let [msg (<! response-chan)]
        (do
          (a/put! event-ch (generate-string {:response msg}))
          (recur))
        (a/close! event-ch)))))

(def my-json-interceptor
  {:name  ::my-json-interceptor
   :enter (fn [{:keys [request] :as ctx}]
            (if (#{:post :put} (:request-method request))
              (let [raw-body-str (slurp (:body request))
                    json-params (parse-string raw-body-str true)]
                (assoc-in (assoc-in ctx [:request :json-params] json-params)
                          [:request :raw-body-str] raw-body-str))
              ctx))})

(def routes
  #{["/" :get
     [index-handler]
     :route-name :index]
    ["/api/generate" :post
     [my-json-interceptor (sse/start-event-stream stream-ready)]
     :route-name :stream]})

I think it makes a lot of sense to use Pedestal for async apps considering both core.async and Pedestal are maintained by the same folks (and I believe they use both libraries extensively in-house). That way you can passively benefit from any improvements or updates to core.async over time - plus it seems like there’s a lot of cool stuff you can do with core.async in general. It can be difficult to use, so you are pretty much required to be willing to read the source and ask for help on Clojurians slack when needed.

For a project as small as this, Go was way easier to work with. I’ve written Clojure on and off over the last few years and I still spent 10-20x longer getting those implementations working. With Go, it’s really easy to just pick the first dumb implementation that pops into your head and it will work. And of course, trivial deployment is always a relief.

I do still reach for Clojure from time to time, both because I enjoy the development process of building projects in small pieces without ceremony - i.e., you don’t have to create mini-projects or something similar to experiment with new ideas, you just immediately try them out from whatever file you’re currently working in. I also find it much easier to maintain a bird’s eye view of your codebase I’m working on. Clojure code tends to be much higher level and I find it easier to reason about when I’m focused, provided I’ve kept everything logically organized and somewhat tidy.

A couple of years ago, I worked on an exploratory protoype of web application to get my hands dirty with Clojure and Datomic. I put together a prototype web application with a frontend written using Tailwind CSS and Alpine JS and a backend written using Pedestal and Datomic. I chose to hand-roll everything including authentication and role based access control for various user groups.

From the outside it looks like a typical webapp, but the application internals are unusual and worth discussing.

Brief Overview of Clojure

Clojure is a JVM programming language with some distinguishing features: - dynamic scripting language designed to keep a running process that can be manipulated (similar to Python, but with more interaction) - emphasizes functional programming, but allows for mutation where needed - uses Lisp syntax, i.e. (println "hello") instead of System.out.println("hello")

Clojure has a very solid selection of libraries for web development, lets you develop rapidly by incrementally defining functions without needing to recompile, and cuts away a lot of the boilerplate you find in older enterprise languages. As a result, web programming is one of Clojure’s strongest use cases.

For these types of applications, deployment works pretty much the same as it does for Java; you can package up Clojure applications into an archive that include all dependencies, drop it into a server, and run it. I spun up a $5/mo VPS, installed the JDK and Clojure, set up the Datomic peer library backed by Postgres, and had a functional starting point without much trouble.

Once deployed, Clojure applications can listen on a socket for connections. Developers can connect to the program from their editor and run commands to inspect or manipulate the program. In fact, it is common to develop applications while they are running, without the need for recompilation. For medium to large sized projects, this is significant. Java projects I worked on in the past could take fifteen minutes to compile, and larger projects took longer.

Templating

In languages like Java, you tend to use HTML templating languages like Thymeleaf or JSPs to embed dynamic data from the server into pages. It generally looks something like this:

<!DOCTYPE html>
<html>
<head>
    <title>Sample Page</title>
</head>
<body>
    <div th:text="'Hello, ' + ${user}">Hello, World!</div>
</body>
</html>

The Java template above is stored in a file, parsed, and processed by Java code. These work well and are instantly familiar to all web developers, but it’s easy to wind up with either heavy amounts of duplication or a rat’s nest of nested page includes that are difficult to work with.

Clojure, by contrast, has a popular library named Hiccup to render with Clojure data structures directly. Using Hiccup looks like the following:

[:html
 [:head
  [:title "Sample Page"]]
 [:body
  [:div "Hello, " user]]]

The Clojure definition is source code in the syntax of the language. You can access and manipulate it by directly operating on the underlying data structure - a vector. In templating languages, you have to usually have to use a combination of A) splitting templates into many pages and B) using a weird templating language that is different from the main programming language. In practice, that can get very hairy.

Using hiccup makes it easy to maintain page logic as it gets increasingly complex. A lot of the noise gets boiled out and you can operate on HTML output as data, rather than relying on template mechanisms and String processing. The value proposition mirrors that of Clojure vs traditional imperative languages.

Datomic

Cognitect, the stewards of Clojure, created a database called Datomic that serves as an alternative to SQL. Datomic is built on the principles of immutability and stores data as a series of immutable events. More importantly (for me personally, at least) it provides a lot of extremely useful functionality out of the box that you don’t get in SQL.

As the database schema changes, and when data is added or removed, Datomic keeps a complete accounting of events. This approach contrasts with traditional SQL databases where updates or deletions overwrite the original data. In practice, between careful database migrations and database or server logs, you can sometimes stitch together a true historical record of what happened in a database. But as anyone who has worked as an application developer will know - it is far from easy.

If you’ve maintained a traditional database for an extended period of time, you know that it sees incremental changes over time as enhancements are added. So even if you’re able to look at past query results, it can be very difficult to keep track of the entire state of the database at a previous point of time- what the schema was at that time, what data was in there, and so on. In Datomic, all of this information is forever at your fingertips.

You can query the database for its transaction history, and you’ll get back a traversable Clojure sequence of results the same way you get back any other database result. Here’s one way you can do it:

(defn collect-transactions [conn]
  (let [tx-log (d/log conn)]
    (map (fn [tx]
           {:transaction-id (:t tx)
            :timestamp (get-tx-timestamp (d/db conn) (:t tx))
            :data tx})
         (d/tx-range tx-log nil nil))))

And this is what results look like through the eyes of the Morse tool:

Another key feature of Datomic is its use of a declarative query language similar to Clojure’s syntax, which can simplify the learning curve for those already familiar with Clojure.

Database as a Value

In Datomic, the database is treated as a “value.” This implies that the database, at any point in time, represents a specific state that does not change. When you query the database, you are essentially looking at a snapshot of the data as it existed at a particular moment.

There are a lot of cool consequences to this. One of my favorite is the ability to do speculative translations. What this means is that you can develop experimental features to your application, pass your production database to it, and iterate on your concept with no extra work required.

Another benefit is that it is easy to write tests on the fly using live data. It is easy to explore and manipulate your database without worrying about persisting unwanted data - you can simply create and discard branches from any point of your database as desired.

Modifying Schemas

Datomic allows you to change your schema as needed over time. This means you can add new attributes, change the types of existing attributes, or remove attributes that are no longer necessary without disrupting the existing data. These modifications can be done while the database is live and in use, enabling ongoing development and adjustment to evolving data requirements.

I think this is a big deal, because there are definitely projects where you know ahead of time that you’ll be changing direction quite a lot. In a traditional database, this can result in tech debt piling up fairly quickly, or painting yourself into a corner.

As you might expect, there are constraints on this. For example, if your modified schema invalidates existing data, you will need to retract that data first. But overall, evolving your database is far less terrifying than it is in a SQL backed project.

Thoughts on Datomic

In a traditional web app + database deployment, the webapp and the database are two separate behemoths, and each requires a lot of work to spin up and maintain. In enterprise environments, this involves commonly is spread across multiple people, often on different teams.

Another issue with this setup is that SQL has quite a lot of arcane knowledge you have to lug around in order to execute even relatively simple projects. You need to be able to create tables, understand queries, joins, indexing, migration, and best practices associated with all of them. And when you work with your database, it’s generally separate from the application, either from the command line or with some kind of GUI tool.

Datomic has a bit of an initial learning curve because of its use of Datalog syntax and the fact that it operates quite differently from SQL, but I found it to be much simpler to use than SQL after that. You write your schema by passing in maps that lay out your data, update the schema by passing in new maps, seed the database by passing in maps that conform to the schema, and query the database by passing in Clojure vectors that describe what you’re looking for. You don’t need additional dependencies or ORMs, you just interface with the API directly from within your source code. If you’ve ever had to spend time bug fixing decades old ORMs, or migrating off database dependencies because they’re abandoned and are lighting up security scanners, you’ll understand this is a hugely welcome change.

When your application takes shape, it consists of pure functions in your application layer and pure functions in your UI layer through which you thread data you’re entering or retrieving. It comes together quite beautifully.

In many ways, I think that Datomic is to the traditional database as Clojure is to traditional web applications. The value proposition and tradeoffs for both of them with respect to their traditional alternatives is similar- you get the benefits of functional programming at the cost of added space and performance constraints. In Clojure, this means that you (and the folks who build your web servers and other tooling) avoid many gnarly situations that arise from mutable state. In Datomic, you get to treat the database as an immutable value as described above.

Datomic is a bold departure from SQL, and I have a lot of respect for its choices. You may have noticed a lot of its benefits are still somewhat doable in a traditional database, but it often requires quite a lot of extra work, or extra dependencies that bring additional points of failure. Overall, I think Datomic aligns better to real-world persistance requirements of applications than SQL.

It comes at a cost, of course. I’ve read that performance on Datomic Cloud can be quite bad. I personally used the on-prem “peer” library instead, which worked great for me, even when I threw a few stress tests at it.

Some other significant constraints are that Datomic uses a single process to write, and after 10 billion transactions it apparently becomes unweildly. My intuition is that a large number of small to mid sized applications would stay well within its performance budget. The auditing you get for free would already have been such a life saver for a number of projects I’ve worked on in the past.

I ran into a handful of quirks and frustrating moments while building my prototype, mainly stemming from Clojure’s stack traces and error messages from Datomic that were hard to read. Since the language is not very widely used, it can also be much more difficult to get help with specific issues, because you’re less likely to find online posts from other people who ran into them.

Another big drawback is that while the database tracks its “value” over time, you cannot go back and add previous historical data. It only records data as of the moment it passes through the transactor. This presents challenges for situations that will occur often enough in the real world that it bears mentioning. This is called “bi-temporal history”, which is a feature of other similar database systems, but is not a feature of Datomic.

Overall, I really loved working with Datomic. If it fits your problem, it removes so many issues that arise with SQL. I’ll be keeping an eye on the library and looking for opportunities to use it in the future.

Meta Quest headsets let you write 3D programs in C or C++, using OpenGL ES or Vulkan as a rendering backend, and deploy them using the XR Mobile SDK.

I bought a Quest 3 at launch, and I’ve benefited a lot from studying the SDK and experimenting with the sample programs:

1. It gave me a thorough understanding of the OpenXR specification, from its design intent to the low-level details of how it is implemented on the hardware.
2. It significantly improved my understanding of OpenGL ES. I’ve skirted by with much less knowledge when writing other 3D programs in the past, but it was impossible for me to modify the SDK examples and get them to actually run correctly without filling in many knowledge gaps.
3. I learned a lot about how to structure programs in C and C-flavored C++ from the simple & readable SDK example programs, which were written by extremely talented graphics programmers.
4. It was a lot of fun! The best projects are the ones you’re motivated to work on, and bringing your programs to life in VR is such a cool experience.

Given how powerful these devices are, and how much fun they are to hack on, it’s kind of crazy to me how little community or discourse there is online surrounding them.

In this blog post I’ll introduce how developing on Quest devices works. I’ll touch on the inner workings of the SDK, cover the structure of sample projects, and go over a small set of changes I made to one of them.

Meta XR Mobile SDK

The SDK includes a collection of small programs, each consisting of an NDK project that focuses on a particular capability of the device.

In total, you’ll find 18 such projects in the SDK.

For people who are curious about VR development, enjoy low-level programming, and have at least glanced at the core topics covered by OpenGL tutorials that are floating around online (1,2), these devices are a slam dunk for hands-on learning.

Meta Quest = Android + Beefy GPU + Sensors

Meta Quest devices run a modified version of Android operating system. I think it’s helpful to build an understanding of how graphics work in the Android ecosystem at large to gain a deeper understanding of how Quest devices work.

There are three ways you can draw to the screen in Android-land: with the Canvas API, OpenGL ES, and Vulkan.

The Canvas API is built upon the Skia Graphics Library, a cross-platform rendering engine. Android programs that use the Canvas API typically use the official SDK or a cross-platform toolkit like Flutter.

Such toolkits offer pre-built UI components that facilitate basic 2D rendering. They offer solutions for simple graphics, text rendering, and basic animations.

Instead of using the Canvas, you can opt to draw by directly calling into the OpenGL ES and Vulkan runtimes on the device. If you want this level of control over drawing, you probably want to use the “Native Development Kit” or NDK.

Unlike with the traditional SDK, the NDK has you write application code in C or C++ and bridge to it with a very small amount of boilerplate; basically a single Java entry point and a handful of config files. This is the approach taken by the Mobile XR SDK.

Minimal Complexity

The example projects make minimal use of C++ language features. You won’t find elaborate abstractions, excessive indirection, or extensive use of templates.

For each sample project, you get the impression that the developer knew exactly what they wanted to make and just wrote out a common sense implementation of that thing. Older projects are largely self-contained, while later ones started to factor out re-used code into libraries external to their respective projects.

The resulting code mostly takes the form of structs and functions and that favor clarity and readability. You can tell what things do by looking at them.

A benefit of this readability is that you end up build a lot of intuition on “how VR works” by making superficial observations. Take for example the following code from the main render loop of one of the samples.

typedef struct XrPosef {
    XrQuaternionf    orientation;
    XrVector3f       position;
} XrPosef;

// ...

XrPosef xfLocalFromEye[NUM_EYES];

for (int eye = 0; eye < NUM_EYES; eye++) {
    // LOG_POSE( "viewTransform", &projectionInfo.projections[eye].viewTransform );
    XrPosef xfHeadFromEye = projections[eye].pose;
    XrPosef_Multiply(&xfLocalFromEye[eye], &xfLocalFromHead, &xfHeadFromEye);

    XrPosef xfEyeFromLocal = XrPosef_Inverse(xfLocalFromEye[eye]);

    XrMatrix4x4f viewMat{};
    XrMatrix4x4f_CreateFromRigidTransform(&viewMat, &xfEyeFromLocal);

    const XrFovf fov = projections[eye].fov;
    XrMatrix4x4f projMat;
    XrMatrix4x4f_CreateProjectionFov(&projMat, GRAPHICS_OPENGL_ES, fov, 0.1f, 0.0f);

    frameIn.View[eye] = OvrFromXr(viewMat);
    frameIn.Proj[eye] = OvrFromXr(projMat);
}

It becomes immediately clear when reading the above that the position for each eye is independently tracked, and as you might guess, independent frames are rendered for each eye.

You can verify this by looking at that project’s vertex shader (seen below) and note that each eye gets its own matrix transforms:

#define NUM_VIEWS 2
#define VIEW_ID gl_ViewID_OVR
#extension GL_OVR_multiview2 : require
layout(num_views=NUM_VIEWS) in;
in vec3 vertexPosition;
in vec4 vertexColor;
uniform mat4 ModelMatrix;
uniform vec4 ColorScale;
uniform vec4 ColorBias;
uniform SceneMatrices
{
    uniform mat4 ViewMatrix[NUM_VIEWS];
    uniform mat4 ProjectionMatrix[NUM_VIEWS];
} sm;
out vec4 fragmentColor;
void main()
{
    gl_Position = sm.ProjectionMatrix[VIEW_ID] * ( sm.ViewMatrix[VIEW_ID] * ( ModelMatrix * vec4( vertexPosition, 1.0 ) ) );
    fragmentColor = vertexColor * ColorScale + ColorBias;
}

In hindsight, it should be obvious that this is something you do for VR, but I hadn’t really thought about it. I find that the code conveys the concept more clearly and completely than words.

Structure of Sample Projects

Below is an overview of some representative files of a sample project. To build one of the same projects, you can grab a copy of Android Studio Bumblebee (2021.1.1) or earlier, and open settings.gradle. I ran projects in that IDE and worked on them in a separate text editor.

XrVirtualKeyboard
├── Projects
│   └── Android
│       ├── AndroidManifest.xml
│       ├── build
│       ├── build.bat
│       ├── build.gradle
│       ├── build.py
│       ├── jni
│       │   ├── Android.mk
│       │   └── Application.mk
│       ├── local.properties
│       └── settings.gradle
├── Src
│   ├── VirtualKeyboardModelRenderer.cpp
│   ├── VirtualKeyboardModelRenderer.h
│   ├── XrHandHelper.h
│   ├── XrHelper.h
│   ├── XrRenderModelHelper.cpp
│   ├── XrRenderModelHelper.h
│   ├── XrVirtualKeyboardHelper.cpp
│   ├── XrVirtualKeyboardHelper.h
│   └── main.cpp
├── assets
│   └── panel.ktx
├── java
│   └── MainActivity.java
└── res
    └── values
        └── strings.xml

Android.mk contains the Makefile that you’ll be the most interested in. When you add source code or libraries, you’ll be changing this file. They generally look like this:

LOCAL_PATH := $(call my-dir)
include $(CLEAR_VARS)

LOCAL_MODULE := xrvirtualkeyboard

include ../../../../cflags.mk

LOCAL_C_INCLUDES := \
                    $(LOCAL_PATH)/../../../../../SampleCommon/Src \
                    $(LOCAL_PATH)/../../../../../SampleXrFramework/Src \
                    $(LOCAL_PATH)/../../../../../1stParty/OVR/Include \
                    $(LOCAL_PATH)/../../../../../1stParty/utilities/include \
                    $(LOCAL_PATH)/../../../../../3rdParty/stb/src \
                    $(LOCAL_PATH)/../../../../../3rdParty/khronos/openxr/OpenXR-SDK/include \
                    $(LOCAL_PATH)/../../../../../3rdParty/khronos/openxr/OpenXR-SDK/src/common

#
LOCAL_SRC_FILES     :=  ../../../Src/main.cpp \
                        ../../../Src/VirtualKeyboardModelRenderer.cpp \
                        ../../../Src/XrRenderModelHelper.cpp \
                        ../../../Src/XrVirtualKeyboardHelper.cpp \

# include default libraries
LOCAL_LDLIBS            := -llog -landroid -lGLESv3 -lEGL
LOCAL_STATIC_LIBRARIES  := samplexrframework
LOCAL_SHARED_LIBRARIES  := openxr_loader

include $(BUILD_SHARED_LIBRARY)

$(call import-module,SampleXrFramework/Projects/Android/jni)
$(call import-module,OpenXR/Projects/AndroidPrebuilt/jni)

The entry point of each project contains OpenGL and OXR initialization before entering the event loop.

Graphics initialization is done with EGL in the same manner as on other embedded systems (or the desktop):

struct ovrEgl {
    void Clear();
    void CreateContext(const ovrEgl* shareEgl);
    void DestroyContext();
#if defined(XR_USE_GRAPHICS_API_OPENGL_ES)
    EGLint MajorVersion;
    EGLint MinorVersion;
    EGLDisplay Display;
    EGLConfig Config;
    EGLSurface TinySurface;
    EGLSurface MainSurface;
    EGLContext Context;
#elif defined(XR_USE_GRAPHICS_API_OPENGL)
    HDC hDC;
    HGLRC hGLRC;
#endif //
};

// ...

#if defined(XR_USE_GRAPHICS_API_OPENGL_ES)
void ovrEgl::CreateContext(const ovrEgl* shareEgl) {
    if (Display != 0) {
        return;
    }

    Display = eglGetDisplay(EGL_DEFAULT_DISPLAY);
    ALOGV("        eglInitialize( Display, &MajorVersion, &MinorVersion )");
    eglInitialize(Display, &MajorVersion, &MinorVersion);
    // Do NOT use eglChooseConfig, because the Android EGL code pushes in multisample
    // flags in eglChooseConfig if the user has selected the "force 4x MSAA" option in
    // settings, and that is completely wasted for our warp target.
    const int MAX_CONFIGS = 1024;
    EGLConfig configs[MAX_CONFIGS];
    EGLint numConfigs = 0;
    if (eglGetConfigs(Display, configs, MAX_CONFIGS, &numConfigs) == EGL_FALSE) {
        ALOGE("        eglGetConfigs() failed: %s", EglErrorString(eglGetError()));
        return;
    }
    // blah blah

Logging is done with a macro that delegates to NDK’s provided logging facilities when the code is running on Android, otherwise it just falls back on printf.

ALOGV("Creating passthrough layer");

And OpenGL function calls use the common idiom of being wrapped in a macro that executes the OpenGL function, checks the global error message after execution, and prints an error message if an error flag is set after the line executes.

GL(glBindBuffer(GL_ARRAY_BUFFER, VertexBuffer));

Code that calls into OXR functionality work similar to OpenGL in a few ways.

OXR is an error checking macro that behaves in the same manner as GL.

And like with OpenGL, functions from the OpenXR specification are are dynamically linked to at runtime. This means you won’t find their implementation in the SDK; they are running in their own address space outside of the program and are available to all applications on the device.

OXR(xrLocateViews(
    app.Session,
    &projectionInfo,
    &viewState,
    projectionCapacityInput,
    &projectionCountOutput,
    projections));

Making Modifications to Sample Projects

Some codebases can be very daunting, and it can take a long time before you feel comfortable making changes.

In this case, after skimming the code for a few of the samples, I was able to dive in and start poking at stuff.

One of the first things I did was to modify the spatial anchors project shown above to change the dimensions of the anchors and to draw a texture on one side of the cube.

Here are what the changes to make this work look like:

The dimensions and colors were single line changes in the event loop.

- persistedCube.Model *= OVR::Matrix4f::Scaling(0.01f, 0.01f, 0.05f);
+ persistedCube.Model *= OVR::Matrix4f::Scaling(0.05f, 0.05f, 0.005f);

- persistedCube.ColorBias = OVR::Vector4f(1, 0.5, 0, 1); // Orange
+ persistedCube.ColorBias = OVR::Vector4f(0.3, 0.2, 0.1, 1); // Brown

I dropped stb_image.h into the project and added some code that loads textures from images on the device and maintains a pool of the loaded textures which is cycled through during rendering.

#include "TexturePool.h"
#include <android/log.h>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#include <android/asset_manager_jni.h>
#include <android/asset_manager.h>

std::vector<std::string> texture_filenames = {"NASA1.jpg", "NASA2.jpg", "NASA3.jpg", "NASA4.jpg", "NASA5.jpg", "NASA6.jpg", "NASA7.jpg", "NASA8.jpg"  };

size_t texture_pointer = 0;

std::vector<GLuint> texture_ids;

void texture_pool_init(std::vector<std::string> filenames, AAssetManager *amgr) {
	for (std::string filename : filenames) {
        GLuint texture_id;

        GL(glGenTextures(1, &texture_id));
        texture_ids.push_back(texture_id);
        GL(glBindTexture(GL_TEXTURE_2D, texture_id));

        GL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER,
                           GL_LINEAR_MIPMAP_LINEAR));
        GL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR));
        GL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT));
        GL(glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT));
        GLenum format;

		AAsset* asset = AAssetManager_open(amgr, filename.c_str(), AASSET_MODE_BUFFER);
		if (asset == NULL) {
			ALOGE("Failed to open %s", filename.c_str());
		} else {
			ALOGV("opened file %s", filename.c_str());
		}

		int width, height, channels;
        stbi_set_flip_vertically_on_load(true);
		void *data = stbi_load_from_memory((unsigned char*)AAsset_getBuffer(asset), AAsset_getLength(asset), &width, &height, &channels, 0);
		if (data == NULL) {
			ALOGE("Failed to load %s", filename.c_str());
		} else {
			ALOGV("loaded file %s", filename.c_str());
		}

        if (channels == 3) {
            format = GL_RGB;
        } else if (channels == 4) {
            format = GL_RGBA;
        } else {
            fprintf(stderr, "Unsupported number of channels: %d\n", channels);
            exit(EXIT_FAILURE);
        }

        GL(glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, format, GL_UNSIGNED_BYTE,
                        data));
        GL(glGenerateMipmap(GL_TEXTURE_2D));

        stbi_image_free(data);
	}
}

// reset at the start of each frame to keep the same rendering order
void reset_texture_pointer() {
    ALOGV("resetting texture from %u to 0",texture_pointer);
    texture_pointer = 0;
}

// if more textures are requested than are available, wrap around
GLuint next_texture() {
    size_t last_texture_idx = texture_pointer;
    size_t next_texture_idx = (texture_pointer + 1) % texture_ids.size();
    texture_pointer = next_texture_idx;
    return texture_ids[last_texture_idx];
}

I added UVs to the stored vertices, and split out the indices into two groups: one group to be rendered with the a set of shaders that use colors, and another group to be rendered using a set of shaders that use textures.

void ovrGeometry::CreateCube()
{
    struct ovrCubeVertices
    {
        signed char positions[8][4];
        unsigned char colors[8][4];
        float uvs[8][2];
    };

    static const ovrCubeVertices cubeVertices = {
        // positions
        {
            {-127, -127, -127, +127},
            {+127, -127, -127, +127},
            {-127, +127, -127, +127},
            {+127, +127, -127, +127},
            {-127, -127, +127, +127},
            {+127, -127, +127, +127},
            {-127, +127, +127, +127},
            {+127, +127, +127, +127}},
        // colors
        {
            {0x00, 0x00, 0x00, 0xff},
            {0xff, 0x00, 0x00, 0xff},
            {0x00, 0xff, 0x00, 0xff},
            {0xff, 0xff, 0x00, 0xff},
            {0x00, 0x00, 0xff, 0xff},
            {0xff, 0x00, 0xff, 0xff},
            {0x00, 0xff, 0xff, 0xff},
            {0xff, 0xff, 0xff, 0xff}},
        // uvs
        {
            {0.0f, 0.0f}, // Bottom back left corner
            {1.0f, 0.0f}, // Bottom back right corner
            {0.0f, 1.0f}, // Top back left corner
            {1.0f, 1.0f}, // Top back right corner
            {0.0f, 0.0f}, // Bottom front left corner
            {1.0f, 0.0f}, // Bottom front right corner
            {0.0f, 1.0f}, // Top front left corner
            {1.0f, 1.0f}, // Top front right corner
        }
    };

    static const unsigned short cube_indices[30] = {
        0, 2, 1, 2, 3, 1, // back
        6, 7, 2, 2, 7, 3, // top
        4, 0, 5, 5, 0, 1, // bottom
        0, 4, 2, 2, 4, 6, // left
        5, 1, 7, 7, 1, 3  // right
    };

    static const unsigned short tex_indices[6] = {
        4, 5, 6, 6, 5, 7, // front
    };

    VertexCount = 8;
    ColorIndexCount = 30;
    TextureIndexCount = 6;
    IndexCount = 0;

    VertexAttribs[0].Index = VERTEX_ATTRIBUTE_LOCATION_POSITION;
    VertexAttribs[0].Size = 4;
    VertexAttribs[0].Type = GL_BYTE;
    VertexAttribs[0].Normalized = true;
    VertexAttribs[0].Stride = sizeof(cubeVertices.positions[0]);
    VertexAttribs[0].Pointer = (const GLvoid *)offsetof(ovrCubeVertices, positions);

    VertexAttribs[1].Index = VERTEX_ATTRIBUTE_LOCATION_COLOR;
    VertexAttribs[1].Size = 4;
    VertexAttribs[1].Type = GL_UNSIGNED_BYTE;
    VertexAttribs[1].Normalized = true;
    VertexAttribs[1].Stride = sizeof(cubeVertices.colors[0]);
    VertexAttribs[1].Pointer = (const GLvoid *)offsetof(ovrCubeVertices, colors);

    VertexAttribs[2].Index = VERTEX_ATTRIBUTE_LOCATION_UV; // UV attribute
    VertexAttribs[2].Size = 2;
    VertexAttribs[2].Type = GL_FLOAT;
    VertexAttribs[2].Normalized = false;
    VertexAttribs[2].Stride = sizeof(cubeVertices.uvs[0]);
    VertexAttribs[2].Pointer = (const GLvoid *)offsetof(ovrCubeVertices, uvs);

    GL(glGenBuffers(1, &VertexBuffer));
    GL(glBindBuffer(GL_ARRAY_BUFFER, VertexBuffer));
    GL(glBufferData(GL_ARRAY_BUFFER, sizeof(cubeVertices), &cubeVertices, GL_STATIC_DRAW));
    GL(glBindBuffer(GL_ARRAY_BUFFER, 0));

    GL(glGenBuffers(1, &ColorIndexBuffer));
    GL(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, ColorIndexBuffer));
    GL(glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(cube_indices), cube_indices, GL_STATIC_DRAW));
    GL(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0));

    GL(glGenBuffers(1, &TextureIndexBuffer));
    GL(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, TextureIndexBuffer));
    GL(glBufferData(GL_ELEMENT_ARRAY_BUFFER, sizeof(tex_indices), tex_indices, GL_STATIC_DRAW));
    GL(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, 0));
}

I re-used the existing color shaders and added new ones for the texture face:

static const char CUBE_TEX_VERTEX_SHADER[] =
    R"GLSL(
        #define NUM_VIEWS 2
        #define VIEW_ID gl_ViewID_OVR
        #extension GL_OVR_multiview2 : require
        layout(num_views=NUM_VIEWS) in;
        in vec3 vertexPosition;
        in vec2 vertexUv;
        uniform mat4 ModelMatrix;
        uniform SceneMatrices
        {
            uniform mat4 ViewMatrix[NUM_VIEWS];
            uniform mat4 ProjectionMatrix[NUM_VIEWS];
        } sm;
        out vec2 UV;
        void main()
        {
            gl_Position = sm.ProjectionMatrix[VIEW_ID] * ( sm.ViewMatrix[VIEW_ID] * ( ModelMatrix * vec4( vertexPosition, 1.0 ) ) );
            UV = vertexUv;
        }
    )GLSL";

static const char CUBE_TEX_FRAGMENT_SHADER[] =
    R"GLSL(
        in vec2 UV;
        out vec4 color;

        uniform sampler2D Texture0;
        void main()
        {
            color = vec4(texture( Texture0, UV ).rgb, 1.0);
        }
        )GLSL";

The main other substantive change was to add a new section to the frame rendering code to loop over the cubes and draw the texture faces:

    GL(glUseProgram(Scene.TextureProgram.Program));
    GL(glBindBufferBase(
           GL_UNIFORM_BUFFER,
           Scene.TextureProgram.UniformBinding[ovrUniform::Index::SCENE_MATRICES],
           Scene.SceneMatrices));
    if (Scene.TextureProgram.UniformLocation[ovrUniform::Index::VIEW_ID] >=
        0) // NOTE: will not be present when multiview path is enabled.
    {
        GL(glUniform1i(Scene.TextureProgram.UniformLocation[ovrUniform::Index::VIEW_ID], 0));
    }
    for (auto c : Scene.CubeData) {
        GL(glBindVertexArray(Scene.Cube.VertexArrayObject));
        GLint loc = Scene.TextureProgram.UniformLocation[ovrUniform::Index::MODEL_MATRIX];
        if (loc >= 0) {
            GL(glUniformMatrix4fv(loc, 1, GL_TRUE, &c.Model.M[0][0]));
        }

        GL(glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, Scene.Cube.TextureIndexBuffer));
        GL(glActiveTexture(GL_TEXTURE0));

        GL(glBindTexture(GL_TEXTURE_2D, c.KaiTextureID));
        GL(glDrawElements(GL_TRIANGLES, Scene.Cube.TextureIndexCount, GL_UNSIGNED_SHORT, nullptr));
    }

    glBindTexture(GL_TEXTURE_2D, 0);
    GL(glBindVertexArray(0));
    GL(glUseProgram(0));

What I’m trying to illustrate is that the code you’re working with in this SDK looks and behaves very similarly to run of the mill OpenGL rendering code. For the fancy stuff, like per-eye matrix transforms, the sample programs give you a strong foundation to build on rather than you having to implement them from scratch.

In addition to the above changes to rendering spatial anchors, I also had fun messing around with XR controls. You call into the OpenXR bindings to do so, but it feels largely the same as dealing with game controller inputs, although these are more varied and much cooler.

There’s a lot of overlap between the APIs, libraries, build systems used both in these projects and in other OpenGL code that has been authored over the last several decades - tutorials, example projects, open source games, and more. When you read those projects and learn from them, you eventually reach a plateau where it’s unclear what to do next to continue learning. The answer here is just to continue getting exposure to more code and to continue building and thinking critically about what you write. This is the point where developing with Meta Quest becomes an extremely appealing.

I found it personally extremely cool to see how much is accomplished with code that is this simple. The code samples are doing interfacing with numerous sensors and 3D rendering on a tight frame budget, while maintaining at a high frame rate, on constrained hardware, using a less performant graphics backend for demonstrative purposes (using Vulkan would cause a large portion of the surface area of these programs to be dedicated to setup and bookkeeping).

You probably already know if you would enjoy hacking on one. If so, I’d definitely recommend picking one up. Quest 2’s are pretty cheap these days, and you’ve always got the return period if you’re not feeling it.

Introduction

In a recent whitepaper, researchers described a technique to take existing pre-trained text-to-image models and embed new subjects, adding the capability to synthesize photorealistic images of the subject contextualized in the model’s output.

A series of implementations were quickly built, finding their way to the Stable Diffusion web UI project in the form of the sd_dreambooth_extension.

This adds the ability to fine-tune models on images of particular subjects.

Images generated this way come out much better than the baseline model, especially when you start adding additional context.

Unfortunately, the process is resource-intensive, and if your system low VRAM it’s easy for dreambooth to crash with an error like this:

RuntimeError: No executable batch size found, reached zero.

In this blog post, I’ll show a few of the optimizations I used to fix this error and get dreambooth running with on my 8GB Nvidia Geforce 1070. I’ll also show a few of the problems I encountered along the way and how I fixed them.

An arch linux system was used, but these optimizations should work on any OS, possibly with minor modifications.

Optimizations

Use xFormers for image generation

xFormers is a library written by facebook research that improves the speed and memory efficiency of image generation. To install it, stop stable-diffusion-webui if its running and build xformers from source by following these instructions.

Next, you’ll need to add a commandline parameter to enable xformers the next time you start the web ui, like in this line from my webui-user.sh:

export COMMANDLINE_ARGS="--medvram --xformers"

The next time you launch the web ui it should use xFormers for image generation.

You’ll also need to make sure to explicitly use xformer when you’re training inside Dreambooth tab of the web UI under Settings > Advanced > Memory Attention (select xformers)

Tweak dreambooth settings

Dreambooth ran succesffully when I used the following settings in the Dreambooth tab of the web ui:

Use LORA: unchecked
Training Steps Per Image (Epochs): 150
batch size: 1
Learning Rate Scheduler: constant with warmup
Learning Rate: 0.000002
Resolution: 512
Use EMA: unchecked
Use 8bit Adam: checked
Mixed precision: fp16
Memory Attention: xformers
Cache Latents: unchecked

Run Stable Diffusion without a graphical environment

I need to squeeze every last meg of VRAM out of my card, so I usually opt not to run a graphical environment on the machine running stable diffusion. On Linux, this means never starting the X server/Wayland compositor.

To do this, first you’ll need to configure the web ui to work from another machine on your local network.

export COMMANDLINE_ARGS="--medvram --listen --xformers"

You’ll want to disable your display manager or similar mechanism and then reboot to free up as much memory as possible.

Then, using another machine on your network (a wifi-connected smartphone works as well), navigate to [ip address of machine running webui]:[port shown in process output] and control it from there.

Reboot between attempted trainings

I’ve found that when I attempt a training, the follow up will often run OOM even when it should not.

Restarting webui or ensuring that there are no leaking GPU resources with nvtop do not seem to solve the issue.

However, doing a full restart will cause the same configuration that ran OOM will succeed. Until I figure out what is leaking GPU memory like a sieve, I’ll probably continue doing full reboots between training.

Avoiding LORA

Many guides online recommend using LORA to reduce memory usage. However, on my setup, using LORA will cause runs to fail that succeed when I disable it.

If LORA isn’t working for you, try turning it off.

Disabling preview images

Setting Save model frequency, Save preview frequency, and Generate Classification Images Using txt2img gives you preview images at epoch checkpoints that you can view to check for overtraining.

Unfortunately the preview images seem to use GPU memory and cause my trainings to fail.

As a result, I end up only setting Save model frequency. When training completes, I bisect the completed checkpoints and check for overtraining.

For example, if I generate 12 checkpoints, I’ll open up checkpoint 12 and run a few images that include my trained subject in a new context. If the images are incapable of including the new context, e.g. the image above showed Gillian Anderson’s face but didn’t respond to the “starfleet officer” prompt, that means the model has been overtrained.

So I would open checkpoint 6, and if that worked, I would try checkpoint 9, and so on.

Encountered Issues

Below are a hodgepodge of problems I ran into along the way, along with how I overcame them.

Exception training model: ‘Could not run ’xformers::efficient_attention_forward_cutlass’ with arguments from the ‘CUDA’ backend.

I got this when I did not properly follow xformers installation steps.

Name ‘str2optimizer8bit_blockwise’ is not defined

This was caused by python being unable to find a CUDA shared library on my system. To fix it, you’ll need to make sure the folders where CUDA’s libraries are located are set on your system’s path.

On arch linux, I did that by setting environment variable LD_LIBRARY_PATH="/opt/cuda/targets/x86_64-linux/lib:$LD_LIBRARY_PATH".

Img should be PIL image

A number of posts online were saying to check “Apply Horizontal Flip” in the dreambooth settings, but doing so was giving me this error.

I didn’t really look much into it, I just unchecked it and the error no longer occurred.

zipfile.BadZipFile: File is not a zip file

After I initially got dreambooth to train successfully, I exported the result as a .ckpt file. When I tried to load the checkpoint to produce the image, the model wouldn’t load, and I could seezipfile.BadZipFile: File is not a zip file in the server’s output.

I looked at the .ckpt file that was being created and compared it to a valid one from a different model. The valid one looked like a zip archive, but the erronous one I was generating was just a yaml file with a .ckpt filename.

I discovered that setting a custom model name in the “Saving” tab in dreambooth settings caused both a yaml and ckpt to be correctly created, and I could load and use the ckpt file.

Introduction

lchat is a small program that is designed to interface with other programs such as ii that present FIFO interfaces. This blog post will briefly show what that means and how lchat works in practice.

ii

A typical irc client works by negotiating a connection with an IRC server before presenting the user with some kind of interface to interact with the server.

For example, the popular client weechat draws a TUI interface that shows the server you’re connected to, what channels you’ve joined, and so on.

Similar to many other interactive programs, it continually listens for the user’s keyboard input. When the user enters messages or commands, weechat marshals information back and forth between the IRC server, updates its internal state, and redraws its interface.

ii (short for “irc it”) is an IRC client which works at a much lower level. It doesn’t draw anything or listen for keyboard or mouse input at all.

Instead, represents your connection to an IRC server as a hierarchy of files.

Each entity that you interact with in IRC, be it the server (to issue commands), or a channel, or a user, gets its own directory and its own pair of files.

One file is named in, and anything written to that file is sent to the process. Another file is named out, and anything received from the process is written to it.

Reading and writing to these files is not done by ii itself, but by some other existing part of the system, such as echo for writing and cat for reading.

Representing processes by files in this manner is known as using a “Named Pipe” or a “FIFO” in Unix systems.

using ii

If you run ii on its own, you’ll see a concise summary of its usage:

usage: ii -s host [-p port | -u sockname] [-i ircdir]
[-n nickname] [-f fullname] [-k env_pass]

When I run ii -s irc.libera.chat -n eskin, ii connects me to libera.chat. By default this creates a few folders in “~/irc”:

The tty that I ran ii from is now taken over by a process that maintains a connection to the server and displays all messages from the server. To interact with ii, I switch to another terminal and interact with the in and out files created for the server I joined.

Here’s me failing to login:

And here’s me joining a room:

Writing to in files and reading from out files with echo and cat are a little unwieldly, which is where lchat comes in.

lchat

lchat expects a command line argument containing the location of a directory containing in and out files.

The command starts a process that continually reads from the out file and displaying its content in the terminal. At the same time, it continually provides a message area for you to type messages and submit them by pressing ‘Enter’. When you do, the contents of the message area are sent to the in file.

In the image below, I navigate to the #linux subdirectory which is where the in and out files for the linux channel exist on my filesystem and the enter lchat . to start lchat in that location.

lchat will work with any program that uses FIFO this way, including a few programs listed on the lchat homepage.

Installation

If you’d like to try out lchat, you’ll first need to grab ii:

git clone https://git.suckless.org/ii
cd ii

If you want to use SSL, you’ll want to apply a patch to the repository:

wget https://tools.suckless.org/ii/patches/ssl/ii-2.0-ssl.diff
git apply ii-2.0-ssl.diff

Then build and install ii:

make
sudo make install

Finally, you’ll want to clone, build, make, and install lchat:

git clone https://git.suckless.org/lchat
cd lchat
make
sudo make install

I’m a fan of the filebased approach to irc, so that I can keep it passively running without keeping track of an active window.

ii and lchat are fun, hackable projects that are fun to play around with. lchat was the first use I saw of libgrapheme in the wild, so that was cool to see. Check it out if you’re in to that sort of stuff!