DEV Community: katyu

Learn Zig by writing own simple Redis client

katyu — Mon, 27 Jan 2025 00:56:51 +0000

Let's explore Zig together by writing a basic Redis client capable of GET/SET operations. We'll start with Redis protocol basics, set up our environment, then dive into Zig code. No prior Zig experience needed - we're figuring this out together

First things first. The Redis protocol

The whole idea is quite simple:

1. Argument count prefixed with *
2. Each argument's length prefixed with $
3. CRLF (\r\n) line terminators

Which looks like this:

*2\r\n
$3\r\n
GET\r\n
$11\r\n
key_is_here\r\n

Send it to Redis, and if there is a value, it will be your value (with its length first):

$10\r\n
your_value\r\n

Or in case of SET:

*3\r\n
$3\r\n
SET\r\n
$11\r\n
key_is_here\r\n
$10\r\n
your_value\r\n

it will be a simple "OK": +OK\r\n

(If you need more details, read it here if you must https://redis.io/docs/latest/develop/reference/protocol-spec/)

Now we need to get a connection to Redis, form commands and parse responses. Speaking of Redis, let's set up a minimal env to work with.

Once in Docker, always in Docker

Someone may argue, but I always prefer to keep it clean and in Docker. You may have your Redis up and running already, if not, let's do it quickly.

Assuming, you have Docker set up already, have the following Dockerfile

FROM redis:7.2
# Copy custom redis.conf if you want to customize configuration
COPY redis.conf /usr/local/etc/redis/redis.conf
# Modify the redis.conf to allow external connections
RUN echo "bind 0.0.0.0" >> /usr/local/etc/redis/redis.conf
# Expose Redis port
EXPOSE 6379
# Start Redis with the custom config
CMD ["redis-server", "/usr/local/etc/redis/redis.conf"]

with a redis.conf to adjust if needed:

bind 0.0.0.0
protected-mode no
port 6379

Build and run - it should be fast

docker build -t my-redis .
docker run -d -p 6379:6379 --name redis-dev my-redis

If you come from the Python world, you may want to run a simple script to check, that your Redis is really running and accessible:

import redis

r = redis.StrictRedis(host='localhost', port=6379, db=0)
r.set('foo', 'bar')
value = r.get('foo')
print(value)

Trying to get in...

All of that was just to prepare us for the main dish. Now it's time to set up our minimal Zig project.

You can take a look here https://zig.guide/getting-started/installation or here https://ziglang.org/learn/getting-started/ but eventually you should be able to init the project, build it and get:

All your codebase are belong to us.
Run `zig build test` to run the tests.

Don't forget to run tests, take time to explore. Make sure you are comfortable to build and run the minimal init project, making a few changes, rebuild, rerun, etc.

Serious work is starting here

Let's bring in the networking package. Zig's standard library provides everything we need for TCP communication and it's std.net:

const std = @import("std");
const net = std.net;
const mem = std.mem;

Have a closer look here https://github.com/ziglang/zig/blob/master/lib/std/net.zig

With these imports, establishing a connection is straightforward. Two lines of code are all we need:

const address = try net.Address.parseIp(host, port);
const stream = try net.tcpConnectToAddress(address);

What's this stream thing we're getting back? It's Zig's abstraction over network communication, providing methods for reading and writing data. Our Redis struct takes ownership of this stream:

pub const Redis = struct {
    stream: net.Stream,
    allocator: std.mem.Allocator,
};

What's Allocator? Zig's allocators are like memory managers. The std.mem.Allocator type provides flexible memory allocation strategies with page, fixed, etc. Read quickly here: https://zig.guide/standard-library/allocators

We'll use the page allocator because it's simplest/default one (ask your OS for entire pages of memory). Seems it's not the most optimal way for prod-like code. Better read more about it, but for now:

const allocator = std.heap.page_allocator;

You might also noticed the try keyword appearing in our code. This is Zig's error handling mechanism in action. When a function can fail, it returns an error union, and try unwraps it elegantly. If something goes wrong, the error propagates up. But sometimes we want to handle errors explicitly, and that's where catch comes in:

const response_type = reader.readByte() catch |err| {
    return mapErrorCases(err);
};

By the way, you've probably noticed type annotations everywhere. That's because Zig is statically typed (in case you missed it), meaning all types must be known at compile time. Look at this function signature:

pub fn init(allocator: std.mem.Allocator, host: []const u8, port: u16) !Redis

Here, []const u8 represents a string (slice of bytes), u16 is an unsigned 16-bit integer, and the !Redis return type indicates this function returns either a Redis instance or an error.

Prepare requests and parse response

Now let's focus on the main objective - GET/SET commands for Redis.

First, preparing the commands, we may have a simple abstraction to help us:

pub const Command = struct {
    args: std.ArrayList([]const u8),

    pub fn init(allocator: mem.Allocator) Command {
        return .{ .args = std.ArrayList([]const u8).init(allocator) };
    }

    pub fn add(self: *Command, arg: []const u8) !void {
        try self.args.append(arg);
    }

    pub fn encode(self: Command, allocator: mem.Allocator) ![]u8 {
        return encodeCommand(allocator, self.args.items);
    }
};

where arg: []const u8 will be actually our string GET or SET and their args.

Once we have our command, we need to transform it into the Redis protocol format. We encode them like this (you get more context from the code below, when we use it):

pub fn encodeCommand(allocator: mem.Allocator, args: []const []const u8) ![]u8 {
    var list = std.ArrayList(u8).init(allocator);
    defer list.deinit();

    try std.fmt.format(list.writer(), "*{d}\r\n", .{args.len});

    for (args) |arg| {
        try std.fmt.format(list.writer(), "${d}\r\n", .{arg.len});
        try list.appendSlice(arg);
        try list.appendSlice("\r\n");
    }

    return list.toOwnedSlice();
}

When Redis sends back data, we need to read it line by line. Here's our custom line reader that handles the CRLF sequence properly:

pub fn readLine(allocator: mem.Allocator, reader: net.Stream.Reader) ![]const u8 {
    var list = std.ArrayList(u8).init(allocator);
    defer list.deinit();

    while (true) {
        const byte = reader.readByte() catch |err| {
            return mapErrorCases(err);
        };

        if (byte == '\r') {
            const next_byte = try reader.readByte();
            if (next_byte == '\n') {
                break;
            }
            try list.append(byte);
            try list.append(next_byte);
        } else {
            try list.append(byte);
        }
    }

    return list.toOwnedSlice();
}

defer ensures resources are released when we exit the current scope.

Then comes to interpreting Redis responses.

Redis responses start with:

'+' Simple strings
'$' Bulk strings
'-' Errors

Our decoder handles each case accordingly, let's see the code and then explain it in details:

pub fn decodeResponse(allocator: mem.Allocator, reader: net.Stream.Reader) ![]const u8 {
    const response_type = reader.readByte() catch |err| {
        return mapErrorCases(err);
    };

    switch (response_type) {
        '+' => {
            return try readLine(allocator, reader);
        },
        '$' => {
            const line = try readLine(allocator, reader);
            defer allocator.free(line);

            const line_length = std.fmt.parseInt(i64, line, 10) catch {
                return RedisError.InvalidResponse;
            };

            if (line_length == -1) {
                return RedisError.NullResponse;
            }

            const buffer = try allocator.alloc(u8, @intCast(line_length));

            const bytes_read = try reader.readAll(buffer);
            if (bytes_read != line_length) {
                return RedisError.ReadError;
            }

            const cr = try reader.readByte();
            const lf = try reader.readByte();
            if (cr != '\r' or lf != '\n') {
                return RedisError.InvalidResponse;
            }

            return buffer;
        },
        '-' => {
            const err_msg = try readLine(allocator, reader);
            defer allocator.free(err_msg);
            return RedisError.InvalidResponse;
        },
        else => {
            return RedisError.UnknownResponseType;
        },
    }
}

If the response_type is '+', it indicates a simple string response.
The function reads the rest of the line using readLine and returns it.

If the response_type is '$', it indicates a bulk string response.
The function reads the length of the bulk string from the next line.

If the length is -1, there is null response, otherwise, it allocates a buffer of the specified length and reads the bulk string into the buffer.

It then reads the trailing CRLF characters to ensure the response is correctly formatted. It returns the buffer containing the bulk string

And finally, if the response_type is '-', it indicates an error response.

One of Zig's distinctive features is its approach to error handling. The language requires us to handle all possible error cases explicitly. That's why we have this comprehensive error mapping:

pub fn mapErrorCases(err: anyerror) RedisError {
    return switch (err) {
        error.EndOfStream => RedisError.EmptyResponse,
        error.NotOpenForReading => RedisError.ConnectionError,
        error.SystemResources => RedisError.ConnectionError,
        error.WouldBlock => RedisError.ReadError,
        error.ConnectionResetByPeer => RedisError.ConnectionError,
        error.BrokenPipe => RedisError.ConnectionError,
        error.ConnectionTimedOut => RedisError.ConnectionError,
        error.InputOutput => RedisError.ReadError,
        error.Unexpected => RedisError.ReadError,
        error.AccessDenied => RedisError.ReadError,
        error.OperationAborted => RedisError.ReadError,
        error.IsDir => RedisError.ReadError,
        error.SocketNotConnected => RedisError.ConnectionError,
        else => RedisError.InvalidResponse,
    };
}

While this level of error handling might seem excessive coming from other languages, seems it's a fundamental part of Zig's philosophy. Or maybe I got something wrong and best practices slightly different.

The compiler ensures we've thought about every possible failure mode, making our code more reliable and easier to maintain in the long run. Perks of being statically typed.

Put together the client

Now we can implement our Redis client by combining all the components we've created. In Zig, we define our client as a struct with associated methods:

pub const Redis = struct {
    stream: net.Stream,
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator, host: []const u8, port: u16) !Redis {
        const address = try net.Address.parseIp(host, port);
        const stream = try net.tcpConnectToAddress(address);

        return Redis{
            .stream = stream,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Redis) void {
        self.stream.close();
    }

    ...
};

The struct introduces several Zig patterns. First, notice how the init function serves as our constructor, returning a new Redis instance. The self parameter in deinit is similar to this in other languages, representing the instance itself. The asterisk (*) before Redis in the deinit signature indicates we're receiving a pointer to the struct.

Methods that interact with Redis follow a similar pattern:

    pub fn get(self: *Redis, key: []const u8) ![]const u8 {
        var cmd = Command.init(self.allocator);
        try cmd.add("GET");
        try cmd.add(key);

        try self.writeCommand(&cmd);
        return self.readResponse();
    }

    pub fn set(self: *Redis, key: []const u8, value: []const u8) ![]const u8 {
        var cmd = Command.init(self.allocator);
        try cmd.add("SET");
        try cmd.add(key);
        try cmd.add(value);

        try self.writeCommand(&cmd);
        return self.readResponse();
    }

These methods demonstrate Zig's approach to memory and error handling. Each operation that might fail returns an error union (denoted by the ! in the return type). The writeCommand and readResponse helper methods handle the low-level communication:

    fn writeCommand(self: *Redis, cmd: *Command) !void {
        const encoded = try cmd.encode(self.allocator);
        defer self.allocator.free(encoded);
        try self.stream.writer().writeAll(encoded);
    }

    fn readResponse(self: *Redis) ![]const u8 {
        return decodeResponse(self.allocator, self.stream.reader());
    }

The Final Usage

With our Redis client implementation complete, we can create a command-line interface to interact with it. The main function serves as the entry point for our application, handling argument parsing and command execution:

pub fn main() !void {
    const host = "127.0.0.1";
    const port = 6379;

    const allocator = std.heap.page_allocator;

    var args = try std.process.argsWithAllocator(allocator);
    defer args.deinit();

    _ = args.next();

    const command = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient GET <key>\nUsage: rclient SET <key> <value>\n", .{});

    var redis = Redis.init(allocator, host, port) catch |err| {
        switch (err) {
            error.SystemResources => return std.io.getStdOut().writer().print("System resources error for: {s}:{d}\n", .{ host, port }),
            // and many more... Full implementation handles all connection errors - see final code for complete list
        }
    };
    defer redis.deinit();

    ...
}

The orelse operator provides a clean way to handle the case where no command is provided. The defer keyword we saw before, it's for cleanup.

The command execution logic follows:

    if (std.mem.eql(u8, command, "GET")) {
        const key = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient GET <key>\n", .{});

        const response = try redis.get(key);
        try std.io.getStdOut().writer().print("{s}: {s}\n", .{ key, response });
    } else if (std.mem.eql(u8, command, "SET")) {
        const key = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient SET <key> <value>\n", .{});
        const value = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient SET <key> <value>\n", .{});

        const response = try redis.set(key, value);
        try std.io.getStdOut().writer().print("{s}\n", .{response});
    } else {
        try std.io.getStdOut().writer().print("Usage: rclient GET <key>\nUsage: rclient SET <key> <value>\n", .{});
    }

Now we can test our client. After building the project with zig build, we can execute Redis commands:

./rclient SET mykey "myvalue"
OK

./rclient GET mykey
mykey: myvalue

The error handling we implemented earlier ensures that we get meaningful messages when things go wrong:

# When Redis is not running
./rclient GET mykey
Connection refused for: 127.0.0.1:6379

# When using incorrect syntax
./rclient GET
Usage: rclient GET <key>

To Sum Up

We've explored a bit about Redis protocol, set up and env with Docker, created a minimal Zig project and write one-file code for minimal Redis client. Not too bad for a start.

We haven't touched any testing (unit), it would take much more space. But don't ignore it. As you could see from the very init project template, testing is crucial part of Zig's culture/approach/philosophy. Better write test right away. Static types + tests = reliable app.

A quick one would look like:

test "basic SET/GET" {
    var redis = try Redis.init(std.testing.allocator, "127.0.0.1", 6379);
    defer redis.deinit();

    const res = try redis.set("test", "value");
    try std.testing.expectEqualStrings("OK", res);
}

I will continue exploring Zig and maybe even write a few other articles on the way. May even make it interesting enough. And if I missed something or done it in a wrong ineffective way - let me know.

Put all together

Better in git, but for now as is:

const std = @import("std");
const net = std.net;
const mem = std.mem;

// Let's define some errors, Redis specifics to map possible responses after
pub const RedisError = error{
    InvalidResponse,
    EmptyResponse,
    NullResponse,
    UnknownResponseType,
    ConnectionError,
    ReadError,
} || mem.Allocator.Error;

// Looks like we have to map all the errors from errors? There are ways to make it simpler,
// though keeping it explicit for now
pub fn mapErrorCases(err: anyerror) RedisError {
    return switch (err) {
        error.EndOfStream => RedisError.EmptyResponse,
        error.NotOpenForReading => RedisError.ConnectionError,
        error.SystemResources => RedisError.ConnectionError,
        error.WouldBlock => RedisError.ReadError,
        error.ConnectionResetByPeer => RedisError.ConnectionError,
        error.BrokenPipe => RedisError.ConnectionError,
        error.ConnectionTimedOut => RedisError.ConnectionError,
        error.InputOutput => RedisError.ReadError,
        error.Unexpected => RedisError.ReadError,
        error.AccessDenied => RedisError.ReadError,
        error.OperationAborted => RedisError.ReadError,
        error.IsDir => RedisError.ReadError,
        error.SocketNotConnected => RedisError.ConnectionError,
        else => RedisError.InvalidResponse,
    };
}

pub fn readLine(allocator: mem.Allocator, reader: net.Stream.Reader) ![]const u8 {
    // I guess we can use some fixed buffer here,
    // but let's keep it simple and still use dynamic one
    var list = std.ArrayList(u8).init(allocator);
    defer list.deinit();

    while (true) {
        const byte = reader.readByte() catch |err| {
            return mapErrorCases(err);
        };

        if (byte == '\r') {
            const next_byte = try reader.readByte();
            if (next_byte == '\n') {
                break;
            }
            try list.append(byte);
            try list.append(next_byte);
        } else {
            try list.append(byte);
        }
    }

    return list.toOwnedSlice();
}

pub fn decodeResponse(allocator: mem.Allocator, reader: net.Stream.Reader) ![]const u8 {
    const response_type = reader.readByte() catch |err| {
        return mapErrorCases(err);
    };

    switch (response_type) {
        '+' => {
            return try readLine(allocator, reader);
        },
        '$' => {
            const line = try readLine(allocator, reader);
            defer allocator.free(line);

            const line_length = std.fmt.parseInt(i64, line, 10) catch {
                return RedisError.InvalidResponse;
            };

            if (line_length == -1) {
                return RedisError.NullResponse;
            }

            // TODO: Production code should validate response sizes and use streaming reads
            const buffer = try allocator.alloc(u8, @intCast(line_length));

            const bytes_read = try reader.readAll(buffer);
            if (bytes_read != line_length) {
                return RedisError.ReadError;
            }

            const cr = try reader.readByte();
            const lf = try reader.readByte();
            if (cr != '\r' or lf != '\n') {
                return RedisError.InvalidResponse;
            }

            return buffer;
        },
        '-' => {
            const err_msg = try readLine(allocator, reader);
            defer allocator.free(err_msg);
            return RedisError.InvalidResponse;
        },
        else => {
            return RedisError.UnknownResponseType;
        },
    }
}

pub fn encodeCommand(allocator: mem.Allocator, args: []const []const u8) ![]u8 {
    var list = std.ArrayList(u8).init(allocator);
    defer list.deinit();

    try std.fmt.format(list.writer(), "*{d}\r\n", .{args.len});

    for (args) |arg| {
        try std.fmt.format(list.writer(), "${d}\r\n", .{arg.len});
        try list.appendSlice(arg);
        try list.appendSlice("\r\n");
    }

    return list.toOwnedSlice();
}

pub const Command = struct {
    args: std.ArrayList([]const u8),

    pub fn init(allocator: mem.Allocator) Command {
        return .{ .args = std.ArrayList([]const u8).init(allocator) };
    }

    pub fn add(self: *Command, arg: []const u8) !void {
        try self.args.append(arg);
    }

    pub fn encode(self: Command, allocator: mem.Allocator) ![]u8 {
        return encodeCommand(allocator, self.args.items);
    }
};

pub const Redis = struct {
    stream: net.Stream,
    allocator: std.mem.Allocator,

    pub fn init(allocator: std.mem.Allocator, host: []const u8, port: u16) !Redis {
        const address = try net.Address.parseIp(host, port);
        const stream = try net.tcpConnectToAddress(address);

        return Redis{
            .stream = stream,
            .allocator = allocator,
        };
    }

    pub fn deinit(self: *Redis) void {
        // Is there anything else to do here?
        self.stream.close();
    }

    fn writeCommand(self: *Redis, cmd: *Command) !void {
        const encoded = try cmd.encode(self.allocator);
        defer self.allocator.free(encoded);
        try self.stream.writer().writeAll(encoded);
    }

    fn readResponse(self: *Redis) ![]const u8 {
        return decodeResponse(self.allocator, self.stream.reader());
    }

    pub fn get(self: *Redis, key: []const u8) ![]const u8 {
        var cmd = Command.init(self.allocator);
        try cmd.add("GET");
        try cmd.add(key);

        try self.writeCommand(&cmd);
        return self.readResponse();
    }

    pub fn set(self: *Redis, key: []const u8, value: []const u8) ![]const u8 {
        var cmd = Command.init(self.allocator);
        try cmd.add("SET");
        try cmd.add(key);
        try cmd.add(value);

        try self.writeCommand(&cmd);
        return self.readResponse();
    }
};

// Here is our console demo app. Looks a bit ugly, feels like there are easier way to handle input, but it's a good start.
pub fn main() !void {
    // Keep it simple, connect only to local one and to default port
    const host = "127.0.0.1";
    const port = 6379;

    const allocator = std.heap.page_allocator;

    var args = try std.process.argsWithAllocator(allocator);
    defer args.deinit();

    _ = args.next();

    // Love this "orelse", otherwise there is so much boilerplate to handle input
    const command = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient GET <key>\nUsage: rclient SET <key> <value>\n", .{});

    // So we can get a nice "Connection refused for: 127.0.0.1:6379". You can check by shutting down your Redis instance.
    var redis = Redis.init(allocator, host, port) catch |err| {
        switch (err) {
            error.SystemResources => return std.io.getStdOut().writer().print("System resources error for: {s}:{d}\n", .{ host, port }),
            error.WouldBlock => return std.io.getStdOut().writer().print("Operation would block for: {s}:{d}\n", .{ host, port }),
            error.ConnectionResetByPeer => return std.io.getStdOut().writer().print("Connection reset by peer for: {s}:{d}\n", .{ host, port }),
            error.Unexpected => return std.io.getStdOut().writer().print("Unexpected error for: {s}:{d}\n", .{ host, port }),
            error.InvalidIPAddressFormat => return std.io.getStdOut().writer().print("Invalid IP address format for: {s}:{d}\n", .{ host, port }),
            error.PermissionDenied => return std.io.getStdOut().writer().print("Permission denied for: {s}:{d}\n", .{ host, port }),
            error.AddressFamilyNotSupported => return std.io.getStdOut().writer().print("Address family not supported for: {s}:{d}\n", .{ host, port }),
            error.ProtocolFamilyNotAvailable => return std.io.getStdOut().writer().print("Protocol family not available for: {s}:{d}\n", .{ host, port }),
            error.ProcessFdQuotaExceeded => return std.io.getStdOut().writer().print("Process file descriptor quota exceeded for: {s}:{d}\n", .{ host, port }),
            error.SystemFdQuotaExceeded => return std.io.getStdOut().writer().print("System file descriptor quota exceeded for: {s}:{d}\n", .{ host, port }),
            error.ProtocolNotSupported => return std.io.getStdOut().writer().print("Protocol not supported for: {s}:{d}\n", .{ host, port }),
            error.SocketTypeNotSupported => return std.io.getStdOut().writer().print("Socket type not supported for: {s}:{d}\n", .{ host, port }),
            error.AddressInUse => return std.io.getStdOut().writer().print("Address in use for: {s}:{d}\n", .{ host, port }),
            error.AddressNotAvailable => return std.io.getStdOut().writer().print("Address not available for: {s}:{d}\n", .{ host, port }),
            error.ConnectionRefused => return std.io.getStdOut().writer().print("Connection refused for: {s}:{d}\n", .{ host, port }),
            error.NetworkUnreachable => return std.io.getStdOut().writer().print("Network unreachable for: {s}:{d}\n", .{ host, port }),
            error.ConnectionTimedOut => return std.io.getStdOut().writer().print("Connection timed out for: {s}:{d}\n", .{ host, port }),
            error.FileNotFound => return std.io.getStdOut().writer().print("File not found for: {s}:{d}\n", .{ host, port }),
            error.ConnectionPending => return std.io.getStdOut().writer().print("Connection pending for: {s}:{d}\n", .{ host, port }),
        }
    };
    defer redis.deinit();

    if (std.mem.eql(u8, command, "GET")) {
        const key = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient GET <key>\n", .{});

        const response = try redis.get(key);
        try std.io.getStdOut().writer().print("{s}: {s}\n", .{ key, response });
    } else if (std.mem.eql(u8, command, "SET")) {
        const key = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient SET <key> <value>\n", .{});
        const value = args.next() orelse return std.io.getStdOut().writer().print("Usage: rclient SET <key> <value>\n", .{});

        const response = try redis.set(key, value);
        try std.io.getStdOut().writer().print("{s}\n", .{response});
    } else {
        try std.io.getStdOut().writer().print("Usage: rclient GET <key>\nUsage: rclient SET <key> <value>\n", .{});
    }
}

Quick and Dirty Document Analysis: Combining GOT-OCR and LLama in Python

katyu — Thu, 09 Jan 2025 03:13:31 +0000

Let's explore a way to do OCR + LLM analysis for an image. Will this be the best way given by an expert with decades of experience? Not really. But it comes from someone who takes a similar approach in real life. Think of this as a weekend project version with practical snippets rather than production-ready code. Let's dig in!

What's our goal here?

We're going to build a simple pipeline that can take an image (or PDF), extract text from it using OCR, and then analyze that text using an LLM to get useful metadata. This could be handy for automatically categorizing documents, analyzing incoming correspondence, or building a smart document management system. We'll do it using some popular open-source tools and keep things relatively straightforward.

And yeah, everything below assumes you're already pretty comfortable with HF transformers. If not, check out https://huggingface.co/docs/transformers/en/quicktour - seems like a solid place to start. Though I never did and just learned from examples. I'll get to it... eventually.

What packages do we need?

We'll use torch and transformers for the heavy lifting, plus pymupdf and rich to make our lives easier with some user-friendly console output (I like the rich, so basically we're using it for fun).

import json
import time
import fitz

import torch
from transformers import AutoModel, AutoTokenizer, pipeline

from rich.console import Console
console = Console()

Prepare the image

First off, what image should we use as input? Since we're using Hugging Face here for the primary job, let's use the first page of their leading web page as our test subject. It's a good candidate with both text and complicated formatting - perfect for putting our OCR through its paces.

For a more realistic solution, let's assume our input is a PDF (because let's face it, that's what you'll probably deal with in the real world). We'll need to convert it to PNG format for our model to process:

INPUT_PDF_FILE = "./data/ocr_hf_main_page.pdf"
OUTPUT_PNG_FILE = "./data/ocr_hf_main_page.png"

doc = fitz.open(INPUT_PDF_FILE)
page = doc.load_page(0)
pixmap = page.get_pixmap(dpi=300)
img = pixmap.tobytes()

with console.status("Converting PDF to PNG...", spinner="monkey"):
    with open(OUTPUT_PNG_FILE, "wb") as f:
        f.write(img)

Do the real OCR here

I've played around with various OCR solutions for this task. Sure, there's tesseract and plenty of other options out there. But for my test case, I got the best results with GOT-OCR2_0 (https://huggingface.co/stepfun-ai/GOT-OCR2_0). So let's jump right in with that:

tokenizer = AutoTokenizer.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    device_map="cuda",
    trust_remote_code=True,
)
model = AutoModel.from_pretrained(
    "ucaslcl/GOT-OCR2_0",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    pad_token_id=tokenizer.eos_token_id,
)
model = model.eval().cuda()

What's going on here? Well, default AutoModel and AutoTokenizer, the only special enough part is we're setting up the model to use cuda. And this isn't optional. The model requires CUDA support to run.

Now that we've defined our model, let's actually put it to work on our saved file. Also, we will measure the time and print it. Useful not only to compare with different models, but also to understand if it's even feasible for your use case to wait so long (although it's very quick for our case):

def run_ocr_for_file(func: callable, text: str):
    start_time = time.time()
    res = func()
    final_time = time.time() - start_time

    console.rule(f"[bold red] {text} [/bold red]")
    console.print(res)
    console.rule(f"Time: {final_time} seconds")

    return res

result_text = None
with console.status(
    "Running OCR for the result file...",
    spinner="monkey",
):
    result_text = run_ocr_for_file(
        lambda: model.chat(
            tokenizer,
            OUTPUT_PNG_FILE,
            ocr_type="ocr",
        ),
        "plain texts OCR",
    )

And here's what we get from our original image:

Hugging Face- The Al community building the future.  https: / / hugging face. co/  Search models, datasets, users. . .  Following 0
All Models Datasets Spaces Papers Collections Community Posts Up votes Likes New Follow your favorite Al creators Refresh List black-
forest- labs· Advancing state- of- the- art image generation Follow stability a i· Sharing open- source image generation models
Follow bria a i· Specializing in advanced image editing models Follow Trending last 7 days All Models Datasets Spaces deep see k- a
i/ Deep Seek- V 3 Updated 3 days ago· 40 k· 877 deep see k- a i/ Deep Seek- V 3- Base Updated 3 days ago· 6.34 k· 1.06 k 2.39 k
TRELLIS Q wen/ QV Q- 72 B- Preview 88888888888888888888 888888888888888888 301 Gemini Co der 1 of 3 2025-01-01,9:38 p. m

^ all the text, no formatting, but it's intentional.

GOT-OCR2_0 is pretty flexible - it can output in different formats, including HTML. Here are some other ways you can use it:

# format texts OCR:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='format',
)

# fine-grained OCR:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='ocr',
  ocr_box='',
)
# ... ocr_type='format', ocr_box='')
# ... ocr_type='ocr', ocr_color='')
# ... ocr_type='format', ocr_color='')

# multi-crop OCR:
# ... ocr_type='ocr')
# ... ocr_type='format')

# render the formatted OCR results:
result_text = model.chat(
  tokenizer,
  image_file,
  ocr_type='format',
  render=True,
  save_render_file = './demo.html',
)

Finally try LLM

Now comes the fun part - picking an LLM. There's been endless discussion about which one's best, with articles everywhere you look. But let's keep it simple: what's the LLM everyone and their dog has heard of? Llama. So we'll use Llama-3.2-1B to process out text.

What can we get from the text? Think basic stuff like text classification, sentiment analysis, language detection, etc. Imagine you're building a system to automatically categorize uploaded documents or sort incoming faxes for a pharmacy.

I'll skip the deep dive into prompt engineering (that's a whole other article and I don't believe I will be writing any), but here's the basic idea:

prompt = f"Analyze this text of a document and produce JSON metadata in the following fields: tags, language, confidentiality, priority, category, and summary. Provide ONLY JSON with the first character, no explanations. The following text is the document to analyze: {result_text}"

model_id = "meta-llama/Llama-3.2-1B-Instruct"
pipe = pipeline(
    "text-generation",
    model=model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
messages = [
    {
        "role": "system",
        "content": "You analyze the provided text and generate the JSON output.",
    },
    {"role": "user", "content": prompt},
]

with console.status(
    "Analyzing the text with LLM...",
    spinner="monkey",
):
    outputs = pipe(
        messages,
        max_new_tokens=2048,
    )

By the way, am I doing something hilariously stupid here with prompt/content? Let me know. Pretty new to the "prompt engineering" and do not take it seriously enough yet.

The model sometimes wraps the result in markdown code blocks, so we need to handle that (if anyone knows a cleaner way, I'm all ears):

result_json_str = outputs[0]["generated_text"][-1]
content_str = result_json_str["content"]

# Pretty messy here with formatting =(
if content_str.startswith("```

json") and content_str.endswith("

```"):
    content_str = content_str[7:-3].strip()
elif content_str.startswith("```

") and content_str.endswith("

```"):
    content_str = content_str[3:-3].strip()
elif content_str.startswith("{") and content_str.endswith("}"):
    pass
else:
    raise ValueError("JSON content not found")

parsed_json = json.loads(content_str)

with console.status(
    "Saving the result to a file...",
    spinner="monkey",
):
    with open("result.json", "w") as f:
        json.dump(parsed_json, f, indent=4)

And here's what we typically get as output:

{
    "tags": [
        "Hugging Face",
        "The Al community",
        "building the future",
        "search models",
        "datasets",
        "users"
    ],
    "language": "en",
    "confidentiality": "public",
    "priority": "low",
    "category": "technology",
    "summary": "The Hugging Face community is building the future through the development of search models, datasets, and users."
}

To sum up

We've built a little pipeline that can take a PDF, extract its text using some pretty good OCR, and then analyze that text using an LLM to get useful metadata. Is it production-ready? Probably not. But it's a solid starting point if you're looking to build something similar. The cool thing is how we combined different open-source tools to create something useful - from PDF handling to OCR to LLM analysis.

You can easily extend this. Maybe add better error handling, support for multiple pages, or try different LLMs. Or maybe hook it up to a document management system. Hope you will. It might be a fun task.

Remember, this is just one way to do it - there are probably dozens of other approaches that might work better for your specific use case. But hopefully, this gives you a good starting point for your own experiments! Or a perfect place to teach me in the comments how it's done.