eng2kab-TB

Seq2Seq Transformer Model

A transformer-based sequence-to-sequence model for machine translation task

Model Overview

• Architecture: Transformer (Encoder-Decoder)
• Positional Encoding: Sinusoidal
• Tokenizer Support: Custom target tokenizer/ Bert source tokenizer
• GPU Acceleration: Automatic device detection

Usage

from transformers import AutoTokenizer
from model import Seq2SeqTransformer, Config
import yaml
import torch

# Load config   
with open("config.yaml", "r") as f:
    config_dict = yaml.safe_load(f)

config = Config(**config_dict)
token_config = {
                "add_special_tokens": config.add_special_tokens,
                "return_tensors": config.return_tensors,
             }

# Auto-detect device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load tokenizers
source_tokenizer = AutoTokenizer.from_pretrained("C00LD3R/eng2kab-TB/source_tokenizer")
target_tokenizer = AutoTokenizer.from_pretrained("C00LD3R/eng2kab-TB/target_tokenizer")

# Initialize model
model = Seq2SeqTransformer(
    num_encoder_layers=config.num_encoder_layers,
    num_decoder_layers=config.num_decoder_layers,
    emb_size=config.emb_size,
    nhead=config.nhead,
    src_vocab_size=config.source_vocab_size,
    tgt_vocab_size=config.target_vocab_size,
    dim_feedforward=config.ffn_hid_dim,
    dropout=0.1,
    token_config=token_config,
    source_tokenizer=source_tokenizer,
    target_tokenizer=target_tokenizer
)
# Load weights
with safe_open("model.safetensors", framework="pt", device=device.type) as f:
    state_dict = {key: f.get_tensor(key) for key in f.keys()}

model.load_state_dict(state_dict)

#using greedy
translation = model.translate(model,"hi")

#using beam
translation = model.translate_with_beam_search(model,"hi",lenght_extend= 0, beam_size= 5, raw= False, length_penalty= 0.6)