SocrateX 1.3.1

Custom OCR library based on Transformer architecture.

SocrateX 🚀

SocrateX is a modular, easy-to-use PyTorch framework for training and deploying custom OCR (Optical Character Recognition) Transformer models.

It allows you to build, train, and run inference on OCR models entirely from scratch, or load pre-trained architectures (like the 159M-parameter Socrate model) seamlessly.

Installation

pip install socratex

Quick Start — Training a Custom OCR Model

With SocrateX, you don't need pre-trained weights to get started. You can define your own architecture, generate synthetic training data on the fly, and start training in just a few lines of code.

import SocrateX as sx
import torch
from torch.utils.data import DataLoader

# 1. Initialize a fresh Tokenizer
tokenizer = sx.init_tokenizer()

# 2. Define your architecture (e.g., ~15M parameters)
config = sx.Config(
    d_model=256,
    nhead=4,
    num_layers=4,
    dim_feedforward=1024,
    pool_height=4
)

# 3. Initialize the model
model = sx.init(config=config, tokenizer=tokenizer, device="cuda")

# 4. Generate some quick synthetic data for training!
sx.generate_silly_training_set(
    source="https://raw.githubusercontent.com/first20hours/google-10000-english/master/google-10000-english-no-swears.txt",
    count=1000,
    output_dir="my_dataset"
)

# 5. Load the dataset
images, labels = sx.load_dataset("my_dataset/labels.csv")
dataset = sx.Makeset(images, labels, tokenizer=tokenizer)
sampler = sx.SmartBatchSampler(labels, batch_size=16)
loader  = DataLoader(dataset, batch_sampler=sampler)

# 6. Train the model
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
criterion = torch.nn.CrossEntropyLoss(ignore_index=tokenizer.token_to_id("<pad>"))

trainer = sx.Trainer(model, loader, optimizer, criterion, device="cuda")
for epoch in range(30):
    loss = trainer.train_epoch()
    print(f"Epoch {epoch} | Loss: {loss:.4f}")

# 7. Run Inference
results = model.predict(["my_dataset/word_0.jpg"], function="generate", max_iter=32)
print("Predicted:", results)

Pre-defined Architectures

SocrateX includes factory presets for different model scales:

• sx.cat(tokenizer) — Large model (159M params, d_model=640, 12 layers)
• sx.rat(tokenizer) — Medium model (~80M params, d_model=512, 8 layers)
• sx.mice(tokenizer) — Small model (~40M params, d_model=384, 6 layers)

Hugging Face Integration

SocrateX plays perfectly with the Hugging Face Hub. Our flagship pre-trained model is hosted on Hugging Face and embeds the SocrateX API directly!

from transformers import AutoModel
import torch

# Load the heavy 159M model (weights are downloaded automatically)
model = AutoModel.from_pretrained("ihatebaselines/Socrate", trust_remote_code=True)

# You can use the entire SocrateX API directly on the HuggingFace model object!
results = model.predict(["receipt.jpg"], function="beam_search", beam_width=4)
print(results)

Available Modules

Module	What it does
sx.Config(...)	Define custom Transformer OCR architectures
sx.init(config, tokenizer)	Build model from scratch
sx.load_tokenizer(path)	Load BPE tokenizer from file
sx.init_tokenizer()	Build a fresh tokenizer
sx.load_dataset(path)	Parse CSV/JSON/TXT datasets
sx.Makeset(images, labels, ...)	Create a PyTorch Dataset
sx.SmartBatchSampler(...)	Optimize padding by sorting batches by length
sx.Trainer(...)	Wrapper for clean training loops
sx.predict(...)	Greedy / Fast / Beam Search inference
sx.generate_silly_training_set()	Generate synthetic text image datasets