README.md

---
language: ro
license: mit
tags:
  - ocr
  - handwritten-text-recognition
  - vision
  - transformer
  - pytorch
  - custom-architecture
model_name: Socrate (SocrateX cat)
library_name: socratex
---

# 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

```python
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.

```python
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

```python
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

```python
# 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

```python
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`.

```python
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 |