File size: 4,229 Bytes
04adc54 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """
SocrateX - Custom OCR library based on Transformer architecture.
This library was developed to facilitate training, inference,
and experimentation with SOCRATE models. Everything is modular and customizable.
Quick Start:
-----------------------------------
import SocrateX as sx
# 1. Build a custom model architecture using sx.Config:
config = sx.Config(
d_model=640,
nhead=10,
num_layers=12,
dim_feedforward=2560,
pool_height=4 # nn.AdaptiveMaxPool2d((pool_height, None))
)
tokenizer = sx.load_tokenizer("ocr_bpe_tokenizer.json")
model = sx.init(config=config, tokenizer=tokenizer)
# 2. Build a dataset (height and max_length go here, not in Config):
train_set = model.make_dataset(images, labels, height=32, max_length=64)
# 3. Train:
model.fit(dataloader, optimizer, criterion, epochs=50)
# 4. Predict (inference params go here):
results = model.predict(
image_paths=["document.jpg"],
function="generate",
temp=0.5,
max_iter=64,
penalty=1.15,
top_k=5
)
"""
from .configuration_socrate import SocrateConfig
from .config import Config
from .model import SOCRATE, cat, rat, mice, ResidualBlock, PositionalEncoding, SocratePool
from .dataset import Makeset, SmartBatchSampler, load_dataset
from .trainer import train, Trainer
from .inference import predict, generate, generate_fast, beam_search, extract_crops_from_image
from .tokenizer import init_tokenizer, SocrateXTokenizer
from .synthetic import generate_silly_training_set, generate_silly_testing_set
__all__ = [
"Config",
"SOCRATE",
"SocrateConfig",
"cat",
"rat",
"mice",
"ResidualBlock",
"PositionalEncoding",
"SocratePool",
"Makeset",
"SmartBatchSampler",
"train",
"Trainer",
"predict",
"generate",
"generate_fast",
"beam_search",
"init_tokenizer",
"load_tokenizer",
"SocrateXTokenizer",
"generate_silly_training_set",
"generate_silly_testing_set",
"load_dataset",
]
def load_tokenizer(path="ocr_bpe_tokenizer.json"):
"""
Alias to easily load a tokenizer from a JSON file.
"""
return SocrateXTokenizer.from_file(path)
def load(model_type="cat", weights=None, tokenizer_path="ocr_bpe_tokenizer.json", device="cuda"):
"""
Automatically loads the desired model and tokenizer.
Returns (model, tokenizer).
"""
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_file(tokenizer_path)
if model_type == "cat":
model = cat(tokenizer=tokenizer, weights=weights if weights else cat.pretrained, device=device)
elif model_type == "rat":
model = rat(tokenizer=tokenizer, weights=weights if weights else rat.pretrained, device=device)
elif model_type == "mice":
model = mice(tokenizer=tokenizer, weights=weights if weights else mice.pretrained, device=device)
else:
raise ValueError(f"Unknown model type: {model_type}")
return model, tokenizer
def init(tokenizer=None, config=None, device="cuda"):
"""
Initializes a SOCRATE model from scratch.
Pass an sx.Config() object to fully control the architecture:
config = sx.Config(d_model=256, nhead=4, num_layers=3, pool_height=4)
model = sx.init(config=config, tokenizer=tokenizer)
If config is None, defaults to the cat (158M) architecture.
"""
if tokenizer is None:
raise ValueError("You must provide a tokenizer (sx.init_tokenizer() or sx.load_tokenizer()).")
if config is None:
# Default: cat architecture
config = Config()
hf_config = SocrateConfig(
d_model=config.d_model,
max_len=config.max_len,
nhead=config.nhead,
dim_feedforward=config.dim_feedforward,
activation=config.activation,
norm_first=config.norm_first,
num_layers=config.num_layers,
vocab_size=tokenizer.get_vocab_size(),
pad_id=tokenizer.token_to_id("<pad>"),
bos_id=tokenizer.token_to_id("<bos>"),
eos_id=tokenizer.token_to_id("<eos>"),
)
model = SOCRATE(hf_config, tokenizer=tokenizer, sx_config=config).to(device)
return model
|