Metadata-Version: 2.4
Name: SocrateX
Version: 1.3.1
Summary: Custom OCR library based on Transformer architecture.
Author: Kenyp
Description-Content-Type: text/markdown
Requires-Dist: torch
Requires-Dist: torchvision
Requires-Dist: tokenizers
Requires-Dist: Pillow
Requires-Dist: numpy
Requires-Dist: tqdm
Requires-Dist: opencv-python
Dynamic: author
Dynamic: description
Dynamic: description-content-type
Dynamic: requires-dist
Dynamic: summary

# 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

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

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

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