""" Kurdish Handwritten Paragraph Recognition - Pre-training Script DenseNet121-Transformer Architecture with Curriculum Learning Pre-trains the model on synthetic paragraph images before fine-tuning on real handwritten paragraphs. Usage: python pretrain.py --data_dir ./data/SyntheticParagraphs_12000 --vocab_path ./vocab.json python pretrain.py --data_dir ./data/SyntheticParagraphs_12000 --vocab_path ./vocab.json --no_curriculum """ import os import glob import time import argparse import json import math import random import numpy as np from PIL import Image from datetime import datetime import torch import torch.nn as nn import torch.optim as optim import torch.utils.data as data import torchvision.transforms as transforms import torchvision.models as models from torchvision.transforms import InterpolationMode from torch.nn import functional as F from torch.amp import autocast, GradScaler from tqdm import tqdm import gc # =============================== # Argument Parser # =============================== def parse_args(): parser = argparse.ArgumentParser( description="Kurdish Handwritten Paragraph Recognition - Pre-training") # Data paths parser.add_argument("--data_dir", type=str, required=True, help="Root directory with Training/ and Validation/ subfolders") parser.add_argument("--vocab_path", type=str, required=True, help="Path to vocabulary JSON file (vocab.json)") # Image dimensions parser.add_argument("--img_height", type=int, default=600) parser.add_argument("--img_width", type=int, default=1235) parser.add_argument("--max_seq_len", type=int, default=555) # Training hyperparameters parser.add_argument("--batch_size", type=int, default=16) parser.add_argument("--num_epochs", type=int, default=80) parser.add_argument("--learning_rate", type=float, default=1e-4) parser.add_argument("--grad_clip", type=float, default=5.0) parser.add_argument("--weight_decay", type=float, default=1e-4) parser.add_argument("--seed", type=int, default=42) # Model architecture parser.add_argument("--hidden_size", type=int, default=256) parser.add_argument("--encoder_layers", type=int, default=3) parser.add_argument("--decoder_layers", type=int, default=6) parser.add_argument("--num_heads", type=int, default=8) parser.add_argument("--ff_dim", type=int, default=2048) parser.add_argument("--dropout", type=float, default=0.3) parser.add_argument("--use_upsample", action="store_true", default=True, help="Enable horizontal upsampling layer (default: True)") parser.add_argument("--no_upsample", action="store_true", help="Disable horizontal upsampling layer") # Teacher forcing parser.add_argument("--tf_noise_rate", type=float, default=0.15, help="Teacher forcing noise rate (default: 0.15)") # Curriculum learning parser.add_argument("--no_curriculum", action="store_true", help="Disable curriculum learning (train on all data from start)") # LR scheduler parser.add_argument("--lr_step_size", type=int, default=15, help="StepLR step size in epochs") parser.add_argument("--lr_gamma", type=float, default=0.5, help="StepLR decay factor") # Early stopping parser.add_argument("--patience", type=int, default=15) # Training options parser.add_argument("--mixed_precision", action="store_true", default=True) parser.add_argument("--no_mixed_precision", action="store_true") parser.add_argument("--no_aug", action="store_true", help="Disable data augmentation") # CER computation parser.add_argument("--cer_every", type=int, default=5, help="Compute train CER every N epochs (0 to disable)") parser.add_argument("--cer_max_samples", type=int, default=256, help="Max samples for train CER computation") # Output parser.add_argument("--output_dir", type=str, default="./output", help="Directory to save model and logs") parser.add_argument("--model_name", type=str, default="pretrained_model", help="Base name for saved model file") return parser.parse_args() # =============================== # Vocabulary Loader # =============================== def load_vocabulary(vocab_path): """Load vocabulary from JSON file.""" with open(vocab_path, "r", encoding="utf-8") as f: vocab_data = json.load(f) if "vocab_list" in vocab_data: char_list = vocab_data["vocab_list"] elif "char_to_idx" in vocab_data: mapping = vocab_data["char_to_idx"] char_list = [None] * len(mapping) for char, idx in mapping.items(): char_list[idx] = char else: raise ValueError("Vocabulary JSON must contain 'vocab_list' or 'char_to_idx'") char_to_idx = {char: idx for idx, char in enumerate(char_list)} idx_to_char = {idx: char for idx, char in enumerate(char_list)} return char_list, char_to_idx, idx_to_char # Special token indices (fixed by convention) PAD_TOKEN = 0 SOS_TOKEN = 1 EOS_TOKEN = 2 # =============================== # Helper Functions # =============================== def tensor_to_text(tensor, idx_to_char): """Convert a tensor of character indices to text.""" if isinstance(tensor, torch.Tensor): tensor = tensor.cpu().tolist() text = "" for idx in tensor: if idx == PAD_TOKEN or idx == SOS_TOKEN: continue if idx == EOS_TOKEN: break if idx in idx_to_char: text += idx_to_char[idx] return text def count_lines_in_text(text): """Count the number of lines in a paragraph text.""" if not text: return 0 return text.count('\n') + 1 # =============================== # Curriculum Learning # =============================== # Default schedule: progressive difficulty over 80 epochs DEFAULT_CURRICULUM = [ (1, 8, 1, 1), # Epochs 1-8: 1 line only (9, 16, 1, 2), # Epochs 9-16: 1-2 lines (17, 28, 2, 3), # Epochs 17-28: 2-3 lines (29, 40, 2, 4), # Epochs 29-40: 2-4 lines (41, 52, 3, 5), # Epochs 41-52: 3-5 lines (53, 64, 3, 6), # Epochs 53-64: 3-6 lines (65, 80, 4, 7), # Epochs 65-80: 4-7 lines (full complexity) ] def categorize_paragraphs_by_lines(data_dir): """Group paragraph samples by their line count.""" categories = {} image_files = [] for ext in ["*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg"]: image_files.extend(glob.glob(os.path.join(data_dir, ext))) image_files.extend(glob.glob(os.path.join(data_dir, ext.upper()))) image_files = sorted(list(set(image_files))) for img_path in image_files: label_path = os.path.splitext(img_path)[0] + ".txt" if not os.path.exists(label_path): continue try: with open(label_path, "r", encoding="utf-8") as f: text = f.read().strip() except Exception: try: with open(label_path, "r", encoding="utf-8-sig") as f: text = f.read().strip() except Exception: continue num_lines = count_lines_in_text(text) if num_lines not in categories: categories[num_lines] = [] categories[num_lines].append((img_path, text)) return categories def get_curriculum_stage(epoch, schedule): """Get the min/max line range for the current epoch.""" for start_epoch, end_epoch, min_lines, max_lines in schedule: if start_epoch <= epoch <= end_epoch: return min_lines, max_lines return 1, 7 def filter_paragraphs_by_lines(categories, min_lines, max_lines): """Filter paragraphs to include only those within the line range.""" filtered = [] for num_lines, paragraphs in categories.items(): if min_lines <= num_lines <= max_lines: filtered.extend(paragraphs) return filtered # =============================== # Dataset # =============================== class KurdishParagraphDataset(data.Dataset): """Dataset for Kurdish handwritten paragraph images.""" def __init__(self, root_dir=None, transform=None, max_samples=None, max_seq_len=555, filtered_data=None, img_height=600, img_width=1235, char_to_idx=None): self.transform = transform self.max_seq_len = max_seq_len self.img_height = img_height self.img_width = img_width self.char_to_idx = char_to_idx if filtered_data is not None: self.data = filtered_data else: self.data = [] image_files = [] for ext in ["*.tif", "*.tiff", "*.png", "*.jpg", "*.jpeg"]: image_files.extend(glob.glob(os.path.join(root_dir, ext))) image_files.extend(glob.glob(os.path.join(root_dir, ext.upper()))) image_files = sorted(list(set(image_files))) for img_path in image_files: label_path = os.path.splitext(img_path)[0] + ".txt" if not os.path.exists(label_path): continue try: with open(label_path, "r", encoding="utf-8") as f: text = f.read().strip() except Exception: try: with open(label_path, "r", encoding="utf-8-sig") as f: text = f.read().strip() except Exception: continue if len(text) > 0: self.data.append((img_path, text)) if max_samples and max_samples < len(self.data): random.shuffle(self.data) self.data = self.data[:max_samples] label = "filtered" if filtered_data else root_dir print(f" Loaded {len(self.data)} paragraph images ({label})") def __len__(self): return len(self.data) def __getitem__(self, idx): img_path, text = self.data[idx] image = Image.open(img_path).convert("RGB") orig_width, orig_height = image.size # Aspect-ratio-preserving resize scale = min(self.img_width / orig_width, self.img_height / orig_height) new_width = int(orig_width * scale) new_height = int(orig_height * scale) image = image.resize((new_width, new_height), Image.Resampling.LANCZOS) # Right-aligned on white canvas (RTL script) canvas = Image.new('RGB', (self.img_width, self.img_height), (255, 255, 255)) x_offset = self.img_width - new_width canvas.paste(image, (x_offset, 0)) if self.transform: canvas = self.transform(canvas) # Encode text to indices indices = ([SOS_TOKEN] + [self.char_to_idx.get(c, self.char_to_idx.get(" ", 0)) for c in text] + [EOS_TOKEN]) if len(indices) > self.max_seq_len: indices = indices[:self.max_seq_len - 1] + [EOS_TOKEN] target = torch.LongTensor(indices) return canvas, target, len(indices), text def collate_fn(batch): """Collate function with padding for variable-length targets.""" batch.sort(key=lambda x: x[2], reverse=True) images, targets, lengths, texts = zip(*batch) images = torch.stack(images, 0) max_length = max(lengths) padded = torch.ones(len(targets), max_length).long() * PAD_TOKEN for i, target in enumerate(targets): padded[i, :lengths[i]] = target[:lengths[i]] return images, padded, torch.LongTensor(lengths), texts # =============================== # Augmentation # =============================== def build_train_transform(): """Training augmentation pipeline for paragraph images.""" class ParagraphTransform: def __call__(self, img): if random.random() < 0.5: img = transforms.ColorJitter(brightness=0.15, contrast=0.15)(img) if random.random() < 0.4: img = transforms.RandomAffine( degrees=2, translate=(0.02, 0.02), scale=(0.98, 1.02), shear=(-3, 3), interpolation=InterpolationMode.BILINEAR, fill=255)(img) if random.random() < 0.2: img = transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))(img) img = transforms.ToTensor()(img) if random.random() < 0.3: noise = torch.randn_like(img) * 0.01 img = torch.clamp(img + noise, 0.0, 1.0) img = transforms.Normalize( (0.485, 0.456, 0.406), (0.229, 0.224, 0.225))(img) return img return ParagraphTransform() def build_eval_transform(): """Evaluation transform (normalisation only).""" return transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # =============================== # Positional Encodings # =============================== class PositionalEncoding2D(nn.Module): """2D sinusoidal positional encoding for visual feature maps.""" def __init__(self, d_model, max_h=100, max_w=300): super().__init__() pe = torch.zeros(max_h, max_w, d_model) d_half = d_model // 2 pos_h = torch.arange(0, max_h, dtype=torch.float).unsqueeze(1) div_h = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half)) pe_h = torch.zeros(max_h, d_half) pe_h[:, 0::2] = torch.sin(pos_h * div_h) pe_h[:, 1::2] = torch.cos(pos_h * div_h) pos_w = torch.arange(0, max_w, dtype=torch.float).unsqueeze(1) div_w = torch.exp(torch.arange(0, d_half, 2).float() * (-math.log(10000.0) / d_half)) pe_w = torch.zeros(max_w, d_half) pe_w[:, 0::2] = torch.sin(pos_w * div_w) pe_w[:, 1::2] = torch.cos(pos_w * div_w) for h in range(max_h): for w in range(max_w): pe[h, w, :d_half] = pe_h[h] pe[h, w, d_half:] = pe_w[w] self.register_buffer('pe', pe) def forward(self, x, height, width): _, seq_len, d_model = x.shape pe_2d = self.pe[:height, :width, :].reshape(height * width, d_model) if seq_len <= pe_2d.size(0): pe_2d = pe_2d[:seq_len] else: pad = torch.zeros(seq_len - pe_2d.size(0), d_model, device=x.device) pe_2d = torch.cat([pe_2d, pad], dim=0) return x + pe_2d.unsqueeze(0) class PositionalEncoding1D(nn.Module): """1D sinusoidal positional encoding for decoder sequences.""" def __init__(self, d_model, max_len=1000): super().__init__() pe = torch.zeros(max_len, d_model) position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] = torch.sin(position * div_term) pe[:, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe.unsqueeze(0)) def forward(self, x): return x + self.pe[:, :x.size(1), :] # =============================== # CNN Feature Extractor # =============================== class CNNFeatureExtractor(nn.Module): """DenseNet-121 backbone with optional horizontal upsampling.""" def __init__(self, output_dim=256, use_upsample=True): super().__init__() densenet = models.densenet121(weights=models.DenseNet121_Weights.DEFAULT) self.features = densenet.features backbone_channels = 1024 if use_upsample: self.upsample = nn.Sequential( nn.ConvTranspose2d(backbone_channels, 512, kernel_size=(1, 4), stride=(1, 2), padding=(0, 1)), nn.BatchNorm2d(512), nn.ReLU(inplace=True)) adapt_in = 512 else: self.upsample = None adapt_in = backbone_channels self.adaptation = nn.Sequential( nn.Conv2d(adapt_in, output_dim, kernel_size=1), nn.BatchNorm2d(output_dim), nn.ReLU(inplace=True)) def forward(self, x): features = F.relu(self.features(x), inplace=True) if self.upsample is not None: features = self.upsample(features) features = self.adaptation(features) b, c, h, w = features.shape return features.view(b, c, h * w).permute(0, 2, 1), h, w # =============================== # Transformer OCR Model # =============================== class TransformerOCRParagraphModel(nn.Module): """ DenseNet121-Transformer for end-to-end paragraph recognition. Architecture: 1. DenseNet-121 CNN + optional horizontal upsample 2. 2D positional encoding + Transformer encoder 3. Transformer decoder with 1D positional encoding 4. Linear output projection """ def __init__(self, vocab_size, hidden_size=256, nhead=8, num_encoder_layers=3, num_decoder_layers=6, dim_feedforward=2048, dropout=0.3, use_upsample=True, max_seq_len=555, tf_noise_rate=0.15): super().__init__() self.max_seq_len = max_seq_len self.vocab_size = vocab_size self.tf_noise_rate = tf_noise_rate self.feature_extractor = CNNFeatureExtractor( output_dim=hidden_size, use_upsample=use_upsample) self.pos_encoder_2d = PositionalEncoding2D(hidden_size) self.pos_decoder_1d = PositionalEncoding1D(hidden_size, max_len=max_seq_len) encoder_layer = nn.TransformerEncoderLayer( d_model=hidden_size, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True) self.transformer_encoder = nn.TransformerEncoder( encoder_layer, num_layers=num_encoder_layers) decoder_layer = nn.TransformerDecoderLayer( d_model=hidden_size, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, batch_first=True) self.transformer_decoder = nn.TransformerDecoder( decoder_layer, num_layers=num_decoder_layers) self.token_embedding = nn.Embedding(vocab_size, hidden_size) self.output_projection = nn.Linear(hidden_size, vocab_size) self.hidden_size = hidden_size nn.init.xavier_uniform_(self.token_embedding.weight) nn.init.xavier_uniform_(self.output_projection.weight) def _generate_square_subsequent_mask(self, sz): mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1) return mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, 0.0) def _add_teacher_forcing_noise(self, tgt_input): """Replace random tokens to build decoder robustness.""" if self.tf_noise_rate <= 0 or not self.training: return tgt_input noise_mask = (torch.rand_like(tgt_input.float()) < self.tf_noise_rate) noise_mask = noise_mask & (tgt_input != PAD_TOKEN) & (tgt_input != SOS_TOKEN) random_tokens = torch.randint(3, self.vocab_size, tgt_input.shape, device=tgt_input.device) return torch.where(noise_mask, random_tokens, tgt_input) def forward(self, src, tgt, tgt_key_padding_mask=None): # Encode memory, feat_h, feat_w = self.feature_extractor(src) memory = self.pos_encoder_2d(memory, feat_h, feat_w) memory = self.transformer_encoder(memory) # Decode with teacher forcing tgt_input = self._add_teacher_forcing_noise(tgt[:, :-1]) tgt_embedded = self.pos_decoder_1d(self.token_embedding(tgt_input)) tgt_mask = self._generate_square_subsequent_mask(tgt_embedded.size(1)).to(src.device) tgt_pad_mask = tgt_key_padding_mask[:, :-1] if tgt_key_padding_mask is not None else None output = self.transformer_decoder( tgt_embedded, memory, tgt_mask=tgt_mask, tgt_key_padding_mask=tgt_pad_mask) return self.output_projection(output) def generate_batch(self, imgs, max_length=None): """Auto-regressive greedy batch generation.""" if max_length is None: max_length = self.max_seq_len self.eval() batch_size = imgs.size(0) with torch.no_grad(): memory, feat_h, feat_w = self.feature_extractor(imgs) memory = self.pos_encoder_2d(memory, feat_h, feat_w) memory = self.transformer_encoder(memory) ys = torch.ones(batch_size, 1).fill_(SOS_TOKEN).long().to(imgs.device) finished = torch.zeros(batch_size, dtype=torch.bool, device=imgs.device) for _ in range(max_length - 1): tgt_embedded = self.pos_decoder_1d(self.token_embedding(ys)) tgt_mask = self._generate_square_subsequent_mask(ys.size(1)).to(imgs.device) out = self.transformer_decoder(tgt_embedded, memory, tgt_mask=tgt_mask) out = self.output_projection(out) next_tokens = out[:, -1].argmax(dim=-1) next_tokens[finished] = PAD_TOKEN ys = torch.cat([ys, next_tokens.unsqueeze(1)], dim=1) finished = finished | (next_tokens == EOS_TOKEN) if finished.all(): break return [tensor_to_text(seq, idx_to_char) for seq in ys] # =============================== # Metrics # =============================== def levenshtein_distance(s1, s2): if len(s1) < len(s2): return levenshtein_distance(s2, s1) if len(s2) == 0: return len(s1) prev = range(len(s2) + 1) for c1 in s1: curr = [prev[0] + 1] for j, c2 in enumerate(s2): curr.append(min(prev[j + 1] + 1, curr[j] + 1, prev[j] + (c1 != c2))) prev = curr return prev[-1] def calculate_cer(preds, targets): total_dist = sum(levenshtein_distance(p, t) for p, t in zip(preds, targets)) total_chars = sum(len(t) for t in targets) return total_dist / max(1, total_chars) def calculate_wer(preds, targets): total_dist = sum(levenshtein_distance(p.split(), t.split()) for p, t in zip(preds, targets)) total_words = sum(len(t.split()) for t in targets) return total_dist / max(1, total_words) def evaluate_cer_batch(model, dataloader, device, idx_to_char, max_samples=None): """Compute CER using batch generation.""" model.eval() all_preds, all_targets = [], [] count = 0 with torch.no_grad(): for images, _, _, texts in tqdm(dataloader, desc="Computing CER"): images = images.to(device) if max_samples and count + images.size(0) > max_samples: images = images[:max_samples - count] texts = texts[:max_samples - count] preds = model.generate_batch(images) all_preds.extend(preds) all_targets.extend(texts) count += len(preds) if max_samples and count >= max_samples: break return calculate_cer(all_preds, all_targets) def evaluate_full(model, dataloader, device, idx_to_char): """Full evaluation returning CER, WER, predictions, and targets.""" model.eval() all_preds, all_targets = [], [] with torch.no_grad(): for images, _, _, texts in tqdm(dataloader, desc="Evaluating"): images = images.to(device) preds = model.generate_batch(images) all_preds.extend(preds) all_targets.extend(texts) cer = calculate_cer(all_preds, all_targets) wer = calculate_wer(all_preds, all_targets) return cer, wer, all_preds, all_targets # =============================== # Early Stopping # =============================== class EarlyStopping: def __init__(self, patience=15): self.patience = patience self.counter = 0 self.best_cer = float('inf') self.early_stop = False def __call__(self, val_cer, model, epoch, path): if val_cer < self.best_cer: self.best_cer = val_cer self.counter = 0 torch.save({ 'epoch': epoch, 'model_state_dict': model.state_dict(), 'val_cer': val_cer }, path) print(f" Model saved (Val CER: {val_cer:.4f})") else: self.counter += 1 print(f" Early stopping: {self.counter}/{self.patience}") if self.counter >= self.patience: self.early_stop = True print(" Early stopping triggered.") def reset(self): self.counter = 0 # =============================== # Training Functions # =============================== def train_epoch(model, dataloader, optimizer, criterion, device, scaler, use_mixed_precision=True, grad_clip=5.0): """Train for one epoch.""" model.train() epoch_loss = 0 for images, targets, _, _ in tqdm(dataloader, desc="Training"): images, targets = images.to(device), targets.to(device) tgt_pad_mask = (targets == PAD_TOKEN).to(device) optimizer.zero_grad() if use_mixed_precision: with autocast(device_type='cuda'): outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask) loss = criterion(outputs.reshape(-1, outputs.shape[-1]), targets[:, 1:].reshape(-1)) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) scaler.step(optimizer) scaler.update() else: outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask) loss = criterion(outputs.reshape(-1, outputs.shape[-1]), targets[:, 1:].reshape(-1)) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip) optimizer.step() epoch_loss += loss.item() return epoch_loss / len(dataloader) def evaluate_loss(model, dataloader, criterion, device, use_mixed_precision=True): """Evaluate model loss.""" model.eval() epoch_loss = 0 with torch.no_grad(): for images, targets, _, _ in dataloader: images, targets = images.to(device), targets.to(device) tgt_pad_mask = (targets == PAD_TOKEN).to(device) if use_mixed_precision: with autocast(device_type='cuda'): outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask) loss = criterion(outputs.reshape(-1, outputs.shape[-1]), targets[:, 1:].reshape(-1)) else: outputs = model(images, targets, tgt_key_padding_mask=tgt_pad_mask) loss = criterion(outputs.reshape(-1, outputs.shape[-1]), targets[:, 1:].reshape(-1)) epoch_loss += loss.item() return epoch_loss / len(dataloader) # =============================== # Main # =============================== def main(): global idx_to_char # Used by generate_batch -> tensor_to_text args = parse_args() # Handle flag conflicts use_upsample = args.use_upsample and not args.no_upsample use_mixed_precision = args.mixed_precision and not args.no_mixed_precision use_curriculum = not args.no_curriculum # Seeds torch.manual_seed(args.seed) random.seed(args.seed) np.random.seed(args.seed) # Device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Device: {device}") if torch.cuda.is_available(): print(f"GPU: {torch.cuda.get_device_name(0)}") # Output directory os.makedirs(args.output_dir, exist_ok=True) # Vocabulary char_list, char_to_idx, idx_to_char = load_vocabulary(args.vocab_path) vocab_size = len(char_list) print(f"Vocabulary size: {vocab_size}") # Data directories train_dir = os.path.join(args.data_dir, "Training") val_dir = os.path.join(args.data_dir, "Validation") # Categorize paragraphs by line count print("\nCategorizing paragraphs by line count...") train_categories = categorize_paragraphs_by_lines(train_dir) val_categories = categorize_paragraphs_by_lines(val_dir) total_train = sum(len(v) for v in train_categories.values()) total_val = sum(len(v) for v in val_categories.values()) print(f" Training: {total_train} paragraphs") print(f" Validation: {total_val} paragraphs") if use_curriculum: print("\n Curriculum schedule:") for s, e, mn, mx in DEFAULT_CURRICULUM: label = f"{mn} line only" if mn == mx else f"{mn}-{mx} lines" print(f" Epochs {s:2d}-{e:2d}: {label}") # Transforms train_transform = build_eval_transform() if args.no_aug else build_train_transform() eval_transform = build_eval_transform() # Dataset common kwargs ds_kwargs = dict( max_seq_len=args.max_seq_len, img_height=args.img_height, img_width=args.img_width, char_to_idx=char_to_idx) # Model print("\nInitializing model...") model = TransformerOCRParagraphModel( vocab_size=vocab_size, hidden_size=args.hidden_size, nhead=args.num_heads, num_encoder_layers=args.encoder_layers, num_decoder_layers=args.decoder_layers, dim_feedforward=args.ff_dim, dropout=args.dropout, use_upsample=use_upsample, max_seq_len=args.max_seq_len, tf_noise_rate=args.tf_noise_rate ).to(device) total_params = sum(p.numel() for p in model.parameters()) print(f" Parameters: {total_params:,}") print(f" Upsample: {'ON' if use_upsample else 'OFF'}") print(f" Curriculum: {'ON' if use_curriculum else 'OFF'}") print(f" Teacher forcing noise: {args.tf_noise_rate * 100:.0f}%") # Optimizer, scheduler, criterion optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay) scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma) criterion = nn.CrossEntropyLoss(ignore_index=PAD_TOKEN) scaler = GradScaler('cuda') if use_mixed_precision else None early_stopping = EarlyStopping(patience=args.patience) best_model_path = os.path.join(args.output_dir, f"{args.model_name}.pth") # Log file log_path = os.path.join(args.output_dir, f"{args.model_name}_LOG_{datetime.now():%Y%m%d_%H%M%S}.txt") log_file = open(log_path, 'w', encoding='utf-8') def log(msg): print(msg) log_file.write(msg + '\n') log_file.flush() log(f"\nPre-training started: {datetime.now():%Y-%m-%d %H:%M:%S}") log(f"Config: {vars(args)}") # Training loop current_min_lines = None current_max_lines = None train_loader = None val_loader = None for epoch in range(1, args.num_epochs + 1): start_time = time.time() # Curriculum stage management if use_curriculum: new_min, new_max = get_curriculum_stage(epoch, DEFAULT_CURRICULUM) if new_min != current_min_lines or new_max != current_max_lines: current_min_lines, current_max_lines = new_min, new_max train_filtered = filter_paragraphs_by_lines( train_categories, current_min_lines, current_max_lines) val_filtered = filter_paragraphs_by_lines( val_categories, current_min_lines, current_max_lines) label = (f"{current_min_lines} line only" if current_min_lines == current_max_lines else f"{current_min_lines}-{current_max_lines} lines") log(f"\n Curriculum stage: {label} " f"(train={len(train_filtered)}, val={len(val_filtered)})") train_dataset = KurdishParagraphDataset( transform=train_transform, filtered_data=train_filtered, **ds_kwargs) val_dataset = KurdishParagraphDataset( transform=eval_transform, filtered_data=val_filtered, **ds_kwargs) train_loader = data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, collate_fn=collate_fn, pin_memory=True) val_loader = data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, collate_fn=collate_fn, pin_memory=True) early_stopping.reset() else: if train_loader is None: all_train = [p for ps in train_categories.values() for p in ps] all_val = [p for ps in val_categories.values() for p in ps] train_dataset = KurdishParagraphDataset( transform=train_transform, filtered_data=all_train, **ds_kwargs) val_dataset = KurdishParagraphDataset( transform=eval_transform, filtered_data=all_val, **ds_kwargs) train_loader = data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=0, collate_fn=collate_fn, pin_memory=True) val_loader = data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, collate_fn=collate_fn, pin_memory=True) # Train train_loss = train_epoch(model, train_loader, optimizer, criterion, device, scaler, use_mixed_precision, args.grad_clip) # Train CER (periodic) train_cer = None if args.cer_every > 0 and epoch % args.cer_every == 0: train_cer = evaluate_cer_batch(model, train_loader, device, idx_to_char, args.cer_max_samples) # Validation val_loss = evaluate_loss(model, val_loader, criterion, device, use_mixed_precision) val_cer = evaluate_cer_batch(model, val_loader, device, idx_to_char) scheduler.step() elapsed = time.time() - start_time mins, secs = divmod(elapsed, 60) # Log cer_str = f", Train CER: {train_cer:.4f}" if train_cer is not None else "" log(f"Epoch {epoch}/{args.num_epochs} ({mins:.0f}m {secs:.0f}s) | " f"Train Loss: {train_loss:.4f}{cer_str} | " f"Val Loss: {val_loss:.4f} | Val CER: {val_cer:.4f}") # Early stopping and model saving early_stopping(val_cer, model, epoch, best_model_path) if early_stopping.early_stop: break gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # Final evaluation on full validation set log(f"\nLoading best model for final evaluation...") ckpt = torch.load(best_model_path, map_location=device) model.load_state_dict(ckpt['model_state_dict']) all_val = [p for ps in val_categories.values() for p in ps] full_val_dataset = KurdishParagraphDataset( transform=eval_transform, filtered_data=all_val, **ds_kwargs) full_val_loader = data.DataLoader( full_val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0, collate_fn=collate_fn, pin_memory=True) final_cer, final_wer, preds, targets = evaluate_full( model, full_val_loader, device, idx_to_char) log(f"\nFinal Validation Results (Full Set, {len(full_val_dataset)} paragraphs):") log(f" CER: {final_cer:.4f}") log(f" WER: {final_wer:.4f}") log(f"\nSample Predictions:") for i in range(min(3, len(preds))): log(f"\n--- Sample {i + 1} ---") log(f"Predicted: {preds[i][:200]}") log(f"Actual: {targets[i][:200]}") log(f"\nPre-training complete: {datetime.now():%Y-%m-%d %H:%M:%S}") log(f"Best model saved to: {best_model_path}") log_file.close() print(f"Log saved to: {log_path}") if __name__ == "__main__": main()