Socrate — OCR Transformer Model

Main model fully coded by me. Parameters: 159,207,935

Socrate is a custom Transformer-based OCR model trained to read printed and handwritten text from images. Built with the SocrateX library — a modular, easy-to-use training framework for OCR.

---

Quick Start — Use the Pre-trained Model

from transformers import AutoModel
from huggingface_hub import hf_hub_download

# Load pre-trained Socrate (159M) directly from HuggingFace
model = AutoModel.from_pretrained("ihatebaselines/Socrate", trust_remote_code=True)

# Load the tokenizer
tok_path  = hf_hub_download("ihatebaselines/Socrate", "ocr_bpe_tokenizer.json")
tokenizer = model.make_tokenizer(tok_path)

# Run OCR on an image
results = model.predict(["your_image.jpg"], function="generate", max_iter=64)
print(results)

---

Quick Start — Build Your Own Custom Model (no pretrained weights)

No need to install SocrateX separately. Everything is built into the model object.

from transformers import AutoModel

# Load model from HuggingFace (only needed to access the class + tokenizer)
model = AutoModel.from_pretrained("ihatebaselines/Socrate", trust_remote_code=True)

# 1. Create your own architecture config
cfg = model.create_config(
    d_model=256,
    nhead=4,
    num_layers=4,
    dim_feedforward=1024,
    pool_height=4
)

# 2. Create a tokenizer
tok = model.make_tokenizer()   # fresh tokenizer
# or load an existing one:
# tok = model.make_tokenizer("ocr_bpe_tokenizer.json")

# 3. Build a brand-new model from your config (no pretrained weights)
my_model = model.new(config=cfg, tokenizer=tok, device="cuda")

total_params = sum(p.numel() for p in my_model.parameters())
print(f"Your model has {total_params:,} parameters")

# 4. Load your dataset
images, labels = my_model.load_data("your_label.csv")

# 5. Build a dataset + DataLoader
import torch
from torch.utils.data import DataLoader

dataset = my_model.make_dataset(images, labels)
loader  = DataLoader(dataset, batch_size=16, shuffle=True)

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

trainer = my_model.make_trainer(loader, optimizer, criterion)
for epoch in range(50):
    loss = trainer.train_epoch()
    print(f"Epoch {epoch} | Loss: {loss:.4f}")

# 7. Run inference
results = my_model.predict(["test.jpg"], function="generate", max_iter=64)
print(results)

---

Full API — All Methods on the Model Object

Once loaded with AutoModel.from_pretrained, the model exposes the entire SocrateX API:

Method	Description
model.predict(images, ...)	Run OCR inference on a list of image paths
model.fit(dataloader, optimizer, criterion, ...)	Train the model for N epochs
model.make_dataset(images, labels, ...)	Create a Makeset dataset object
model.create_config(d_model, nhead, ...)	Create a custom architecture config
model.new(config, tokenizer, device)	Build a new model from scratch with your config
model.make_tokenizer(path=None)	Load or create a BPE tokenizer
model.make_trainer(loader, optimizer, criterion)	Returns a Trainer wired to this model
model.load_data(path)	Load images + labels from CSV/JSON/TXT
model.generate_data(source, count, output_dir, mode)	Generate synthetic training data
model.freeze_encoder()	Freeze CNN + Encoder weights (fine-tuning)
model.unfreeze_encoder()	Unfreeze encoder weights
model.load_parameters(path)	Load weights from a .pt checkpoint
model.summary()	Print total/trainable parameter counts

---

Architecture Presets

Model	Params	Notes
cat	159M	This checkpoint (d_model=640, 12 layers)
rat	~80M	d_model=512, 8 layers
mice	~40M	d_model=384, 6 layers
Custom via create_config	Your choice	Full control

---

create_config — All Parameters

cfg = model.create_config(
    d_model=640,            # Transformer hidden dimension
    nhead=10,               # Multi-head attention heads
    num_layers=12,          # Number of Transformer layers
    dim_feedforward=2560,   # FFN hidden dimension (usually 4 * d_model)
    activation="gelu",      # Activation function
    norm_first=True,        # Pre-LayerNorm (modern Transformers)
    max_len=512,            # Max sequence length
    pool_height=4           # SocratePool height (AdaptiveMaxPool2d)
)

---

Inference Functions

# Simple greedy generation (fast, good quality)
results = model.predict(images, function="generate", max_iter=64, temp=0.7, top_k=5, penalty=1.15)

# Faster generation (less accurate but quicker)
results = model.predict(images, function="generate_fast", max_iter=64)

# Beam search (most accurate)
results = model.predict(images, function="beam_search", max_iter=64, beam_width=4)

---

Generate Synthetic Training Data

model.generate_data(
    source="https://raw.githubusercontent.com/first20hours/google-10000-english/master/google-10000-english-no-swears.txt",
    count=1000,
    output_dir="my_train_data",
    mode="train"   # or "test"
)
images, labels = model.load_data("my_train_data/labels.csv")

---

Using the SocrateX Library Directly

If you prefer to use SocrateX as a standalone library (e.g. in a training script), the API is identical — just import SocrateX as sx.

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

# ─── 1. Tokenizer ────────────────────────────────────────────────────────────
tokenizer = sx.load_tokenizer("ocr_bpe_tokenizer.json")
# or build a fresh one:
# tokenizer = sx.init_tokenizer()

# ─── 2. Config ───────────────────────────────────────────────────────────────
config = sx.Config(
    d_model=640,
    nhead=10,
    num_layers=12,
    dim_feedforward=2560,
    pool_height=4
)

# ─── 3. Model ────────────────────────────────────────────────────────────────
model = sx.init(config=config, tokenizer=tokenizer, device="cuda")
# or use a preset:
# model = sx.cat(tokenizer=tokenizer, weights="previous.pt", device="cuda")
# model = sx.rat(tokenizer=tokenizer, device="cuda")
# model = sx.mice(tokenizer=tokenizer, device="cuda")

# ─── 4. Dataset ──────────────────────────────────────────────────────────────
images, labels = sx.load_dataset("label.csv")
dataset = sx.Makeset(images, labels, tokenizer=tokenizer)
sampler = sx.SmartBatchSampler(labels, batch_size=32)
loader  = DataLoader(dataset, batch_sampler=sampler)

# ─── 5. Train ────────────────────────────────────────────────────────────────
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(50):
    loss = trainer.train_epoch()
    print(f"Epoch {epoch} | Loss: {loss:.4f}")

# ─── 6. Inference ────────────────────────────────────────────────────────────
results = model.predict(
    image_paths=["receipt.jpg", "document.png"],
    function="generate",
    max_iter=64,
    temp=0.5,
    top_k=5,
    penalty=1.15
)
print(results)

# ─── 7. Synthetic Data ───────────────────────────────────────────────────────
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="train_data"
)

SocrateX Module Reference

Module	What it does
sx.Config(...)	Define custom architecture
sx.init(config, tokenizer)	Build model from scratch
sx.cat / sx.rat / sx.mice	Preset-sized models
sx.load_tokenizer(path)	Load BPE tokenizer from file
sx.init_tokenizer()	Build a fresh tokenizer
sx.load_dataset(path)	Load CSV/JSON/TXT → (images, labels)
sx.Makeset(images, labels, ...)	Create a PyTorch Dataset
sx.SmartBatchSampler(labels, batch_size)	Length-sorted batch sampler
sx.Trainer(model, loader, opt, crit)	Training loop wrapper
sx.predict(...)	Run inference on image list
sx.generate(...)	Greedy autoregressive generation
sx.beam_search(...)	Beam search decoding
sx.generate_silly_training_set(...)	Generate synthetic word images