Technical Documentation - Guj-Eng OCR-Free Image Translation

Getting Started

Prerequisites

Before setting up the project, ensure you have the following installed:

# Python 3.8 or higher
python --version

# CUDA-capable GPU (recommended)
nvidia-smi

# Git for cloning the repository
git --version

Installation

# Clone the repository
git clone https://github.com/patelchaitany/Guj_eng.git
cd Guj_eng

# Create virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install transformers datasets sentencepiece pillow wandb tqdm matplotlib seaborn

Project Structure

Guj_eng/
├── train.py                 # Main training script
├── validate.py              # Validation script
├── README.md               # Project documentation
├── model/                  # Model implementations
│   ├── transformer.py      # Main Transformer model
│   ├── config.py          # Configuration dataclass
│   ├── encoder.py         # Text encoder
│   ├── decoder.py         # Text decoder
│   ├── visionenc.py       # Vision encoder wrapper
│   └── STN.py             # Spatial Transformer Network
├── dataset/               # Dataset handling
│   ├── dataset.py         # Text translation dataset
│   ├── dataset_image.py   # Image translation datasets
│   ├── ocr_dataset.py     # OCR dataset
│   └── prepare_image.py   # Image preprocessing
├── tokenizer/             # Tokenization
│   └── tokenizer.py       # SentencePiece wrapper
└── learning/              # Learning utilities
    ├── utils.py           # Training utilities
    └── tokenizer.py       # Tokenizer training

Model Architecture

Core Components

The model consists of three main components working together:

1. Vision Encoder (CLIP)

from transformers import CLIPVisionModel

class VisionEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.clip_model = CLIPVisionModel.from_pretrained("openai/clip-vit-base-patch32")
        self.projection = nn.Linear(768, config.comman_embedding_dim)
        
    def forward(self, images):
        # Extract visual features
        vision_outputs = self.clip_model(pixel_values=images)
        image_features = vision_outputs.last_hidden_state
        
        # Project to common embedding space
        projected_features = self.projection(image_features)
        return projected_features

2. Text Encoder

class TextEncoder(nn.Module):
                    def __init__(self, config):
                        super().__init__()
                        self.embedding = nn.Embedding(config.vocab_size, config.comman_embedding_dim)
                        self.transformer_layers = nn.ModuleList([
                            EncoderBlock(config) for _ in range(config.num_encoder_layers)
                        ])
                        self.layer_norm = nn.LayerNorm(config.num_embeddings_encoder)
                        
                    def forward(self, input_ids):
                        # Token embeddings
                        embeddings = self.embedding(input_ids)
                        
                        # Apply transformer layers
                        hidden_states = embeddings
                        for layer in self.transformer_layers:
                            hidden_states = layer(hidden_states)
                            
                        return self.layer_norm(hidden_states)

3. Cross-Attention Decoder

class CrossAttentionDecoder(nn.Module):
                    def __init__(self, config):
                        super().__init__()
                        self.layers = nn.ModuleList([
                            DecoderBlock(config) if i % 2 == 0 else Decoder_Cross(config) 
                            for i in range(config.num_decoder_layers * 2)
                        ])
                        
                    def forward(self, target_ids, encoder_outputs, encoder_mask=None):
                        hidden_states = self.embedding(target_ids)
                        
                        for i, layer in enumerate(self.layers):
                            if i % 2 == 0:  # Self-attention layer
                                hidden_states = layer(hidden_states)
                            else:  # Cross-attention layer
                                hidden_states = layer(hidden_states, encoder_outputs, encoder_mask)
                                
                        return hidden_states

Dataset Preparation

Data Format

The project expects data in specific formats for different tasks:

Text Translation Data

{
  "src": "Hello, how are you?",
  "tgt": "હેલો, તમે કેમ છો?",
  "src_lang": "en",
  "tgt_lang": "gu"
}

Image Translation Data

# Image dataset structure
data/
├── images/
│   ├── img_001.jpg
│   ├── img_002.jpg
│   └── ...
└── annotations/
    ├── train.tsv
    └── val.tsv

# TSV format: image_path \t gujarati_text \t english_text
img_001.jpg	આ એક ઉદાહરણ છે	This is an example

Data Preprocessing

import torchvision.transforms as transforms
from PIL import Image

# Image preprocessing pipeline
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.48145466, 0.4578275, 0.40821073],
        std=[0.26862954, 0.26130258, 0.27577711]
    )  # CLIP normalization
])

# Text preprocessing
def preprocess_text(text, tokenizer, max_length=512):
    tokens = tokenizer.encode(text)
    if len(tokens) > max_length:
        tokens = tokens[:max_length]
    return tokens

Training Process

Multi-Task Training

The model is trained on three tasks simultaneously:

# Training configuration
config = {
    'mt_lambda': 0.5,    # Machine Translation weight
    'tit_lambda': 1.2,   # Text Image Translation weight  
    'ocr_lambda': 0.5,   # OCR weight
    'batch_size': 128,
    'learning_rate': 6e-4,
    'warmup_steps': 1000,
    'max_steps': 10000
}

# Training loop
for step in range(max_steps):
    # Machine Translation Task
    mt_loss = train_mt_batch(model, mt_dataloader, criterion)
    
    # Text Image Translation Task
    tit_loss = train_tit_batch(model, tit_dataloader, criterion)
    
    # OCR Task
    ocr_loss = train_ocr_batch(model, ocr_dataloader, criterion)
    
    # Combined loss
    total_loss = (config['mt_lambda'] * mt_loss + 
                  config['tit_lambda'] * tit_loss + 
                  config['ocr_lambda'] * ocr_loss)
    
    # Backpropagation
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()

Learning Rate Scheduling

def get_lr(step, warmup_steps=1000, max_steps=10000, max_lr=6e-4, min_lr=6e-5):
    if step <= warmup_steps:
        # Linear warmup
        return max_lr * (step / warmup_steps)
    elif step > max_steps:
        return min_lr
    else:
        # Cosine annealing
        decay_ratio = (step - warmup_steps) / (max_steps - warmup_steps)
        return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * decay_ratio))

Usage Examples

Training a Model

# Basic training
python train.py

# Training with custom config
python train.py --batch_size 64 --learning_rate 3e-4 --max_steps 5000

# Resume from checkpoint
python train.py --resume_from checkpoints/transformer_epoch_100.pt

Inference

import torch
from PIL import Image
from model.transformer import Transformer
from model.config import Config
from tokenizer.tokenizer import Tokenizer

# Load model and tokenizer
config = Config(vocab_size=32000)
model = Transformer(config)
model.load_state_dict(torch.load('checkpoints/best_model.pt'))
model.eval()

tokenizer = Tokenizer(["<en|gu>", "<gu|en>", "<en>", "<gu>"])
tokenizer.load('data/tokenizer.model')

# Translate image
def translate_image(image_path, source_lang='gu', target_lang='en'):
    # Load and preprocess image
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0)
    
    # Prepare decoder input
    lang_token = tokenizer.get_token_id(f"<{target_lang}>")
    bos_token = tokenizer.get_token_id("bos")
    decoder_input = torch.tensor([[lang_token, bos_token]])
    
    # Generate translation
    with torch.no_grad():
        for _ in range(50):  # max length
            output = model(decoder_input, image, is_image=True)
            next_token = output[:, -1, :].argmax(dim=-1)
            
            if next_token == tokenizer.get_token_id("eos"):
                break
                
            decoder_input = torch.cat([decoder_input, next_token.unsqueeze(0)], dim=1)
    
    # Decode result
    generated_ids = decoder_input[0, 2:].tolist()
    translation = tokenizer.decode(generated_ids)
    return translation

# Example usage
result = translate_image('path/to/gujarati_image.jpg')
print(f"Translation: {result}")

Evaluation

Validation Script

# Run validation
python validate.py --model_path checkpoints/best_model.pt --data_path data/validation/

# Compute BLEU scores
python validate.py --model_path checkpoints/best_model.pt --metric bleu

# Generate sample translations
python validate.py --model_path checkpoints/best_model.pt --sample_size 100

Metrics

The model is evaluated using standard translation metrics:

BLEU Score: Measures n-gram overlap between prediction and reference
METEOR: Considers synonyms and paraphrases
ChrF: Character-level F-score, good for morphologically rich languages
BERTScore: Semantic similarity using contextual embeddings

Configuration

Model Configuration

@dataclass
class Config:
    vocab_size: int                    # Vocabulary size
    d_model: int = 768                # Model dimension
    nhead: int = 8                    # Number of attention heads
    num_encoder_layers: int = 6       # Encoder layers
    num_decoder_layers: int = 6       # Decoder layers
    dim_feedforward: int = 768        # FFN dimension
    dropout: float = 0.1              # Dropout rate
    max_len: int = 1024              # Maximum sequence length
    num_embeddings_decoder: int = 768 # Decoder embedding dim
    num_embeddings_encoder: int = 768 # Encoder embedding dim
    comman_embedding_dim: int = 512   # Common embedding dim
    control_points: int = 20          # STN control points

Training Configuration

# Hyperparameters
BATCH_SIZE = 128
LEARNING_RATE = 6e-4
WARMUP_STEPS = 1000
MAX_STEPS = 10000
GRADIENT_ACCUMULATION = 32

# Task weights
MT_LAMBDA = 0.5      # Machine Translation
TIT_LAMBDA = 1.2     # Text Image Translation
OCR_LAMBDA = 0.5     # Optical Character Recognition

# Optimization
OPTIMIZER = 'AdamW'
BETAS = (0.9, 0.95)
EPS = 1e-8
WEIGHT_DECAY = 0.01