init

app.py

@@ -0,0 +1,214 @@
+import gradio as gr
+import torch
+import os
+import io
+from gtts import gTTS
+import soundfile as sf
+import tempfile
+import logging
+
+# Import your existing functionality
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+from transformers import Wav2Vec2ForCTC, AutoProcessor
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+
+# Initialize translation model
+checkpoint_dir = "bishaltwr/final_m2m100"  # Change to Hugging Face model ID when deployed
+try:
+    tokenizer = M2M100Tokenizer.from_pretrained(checkpoint_dir)
+    model_m2m = M2M100ForConditionalGeneration.from_pretrained(checkpoint_dir)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model_m2m.to(device)
+    m2m_available = True
+except Exception as e:
+    logging.error(f"Error loading M2M100 model: {e}")
+    m2m_available = False
+
+# Initialize ASR model
+model_id = "bishaltwr/wav2vec2-large-mms-1b-nepali"
+try:
+    processor = AutoProcessor.from_pretrained(model_id)
+    model_asr = Wav2Vec2ForCTC.from_pretrained(model_id, ignore_mismatched_sizes=True)
+    asr_available = True
+except Exception as e:
+    logging.error(f"Error loading ASR model: {e}")
+    asr_available = False
+
+# Initialize X-Transformer model
+try:
+    from inference import translate as xtranslate
+    xtransformer_available = True
+except Exception as e:
+    logging.error(f"Error loading XTransformer model: {e}")
+    xtransformer_available = False
+
+def m2m_translate(text, source_lang, target_lang):
+    """Translation using M2M100 model"""
+    if not m2m_available:
+        return "M2M100 model not available"
+    
+    tokenizer.src_lang = source_lang
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    translated_tokens = model_m2m.generate(
+        **inputs, 
+        forced_bos_token_id=tokenizer.get_lang_id(target_lang)
+    )
+    translated_text = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+    return translated_text
+
+def transcribe_audio(audio_path, language="npi"):
+    """Transcribe audio using ASR model"""
+    if not asr_available:
+        return "ASR model not available"
+    
+    import librosa
+    audio, sr = librosa.load(audio_path, sr=16000)
+    processor.tokenizer.set_target_lang(language)
+    model_asr.load_adapter(language)
+    inputs = processor(audio, sampling_rate=16000, return_tensors="pt")
+    
+    with torch.no_grad():
+        outputs = model_asr(**inputs).logits
+    
+    ids = torch.argmax(outputs, dim=-1)[0]
+    transcription = processor.decode(ids, skip_special_tokens=True)
+    
+    if language == "eng":
+        transcription = transcription.replace('<pad>','').replace('<unk>','')
+    else:
+        transcription = transcription.replace('<pad>',' ').replace('<unk>','')
+    
+    return transcription
+
+def text_to_speech(text):
+    """Convert text to speech using gTTS"""
+    if not text:
+        return None
+    
+    try:
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio:
+            tts = gTTS(text=text)
+            tts.save(temp_audio.name)
+            return temp_audio.name
+    except Exception as e:
+        logging.error(f"TTS error: {e}")
+        return None
+
+def detect_language(text):
+    """Simple language detection function"""
+    english_chars = sum(1 for c in text if c.isascii() and c.isalpha())
+    return "en" if english_chars > len(text) * 0.5 else "ne"
+
+def translate_text(text, model_choice, source_lang=None, target_lang=None):
+    """Main translation function"""
+    if not text:
+        return "Please enter some text to translate"
+        
+    # Auto-detect language if not specified
+    if not source_lang:
+        source_lang = detect_language(text)
+        target_lang = "ne" if source_lang == "en" else "en"
+    
+    # Choose the translation model
+    if model_choice == "XTransformer" and xtransformer_available:
+        return xtranslate(text)
+    elif model_choice == "M2M100" and m2m_available:
+        return m2m_translate(text, source_lang=source_lang, target_lang=target_lang)
+    else:
+        return "Selected model is not available"
+
+# Set up the Gradio interface
+with gr.Blocks(title="Nepali-English Translator") as demo:
+    gr.Markdown("# Nepali-English Translation Service")
+    gr.Markdown("Translate between Nepali and English, transcribe audio, and convert text to speech.")
+    
+    # Set up tabs for different functions
+    with gr.Tabs():
+        # Text Translation Tab
+        with gr.TabItem("Text Translation"):
+            with gr.Row():
+                with gr.Column():
+                    text_input = gr.Textbox(label="Input Text", lines=5)
+                    
+                    with gr.Row():
+                        model_choice = gr.Radio(
+                            choices=["XTransformer", "M2M100"], 
+                            value="XTransformer", 
+                            label="Translation Model"
+                        )
+                        
+                    with gr.Row():
+                        source_lang = gr.Dropdown(
+                            choices=["Auto-detect", "en", "ne"], 
+                            value="Auto-detect", 
+                            label="Source Language",
+                            visible=True
+                        )
+                        target_lang = gr.Dropdown(
+                            choices=["Auto-select", "en", "ne"], 
+                            value="Auto-select", 
+                            label="Target Language",
+                            visible=True
+                        )
+                    
+                    translate_button = gr.Button("Translate")
+                
+                with gr.Column():
+                    translation_output = gr.Textbox(label="Translation Output", lines=5)
+                    tts_button = gr.Button("Convert to Speech")
+                    audio_output = gr.Audio(label="Audio Output")
+        
+        # Speech to Text Tab
+        with gr.TabItem("Speech to Text"):
+            with gr.Column():
+                audio_input = gr.Audio(label="Upload or Record Audio", type="filepath")
+                asr_language = gr.Radio(
+                    choices=["eng", "npi"], 
+                    value="npi", 
+                    label="Speech Language"
+                )
+                transcribe_button = gr.Button("Transcribe")
+                transcription_output = gr.Textbox(label="Transcription Output", lines=3)
+    
+    # Define event handlers
+    def process_translation(text, model, src_lang, tgt_lang):
+        if src_lang == "Auto-detect":
+            src_lang = None
+        if tgt_lang == "Auto-select":
+            tgt_lang = None
+        return translate_text(text, model, src_lang, tgt_lang)
+    
+    def process_tts(text):
+        return text_to_speech(text)
+    
+    def process_transcription(audio_path, language):
+        if not audio_path:
+            return "Please upload or record audio"
+        return transcribe_audio(audio_path, language)
+    
+    # Connect the components
+    translate_button.click(
+        process_translation, 
+        inputs=[text_input, model_choice, source_lang, target_lang], 
+        outputs=translation_output
+    )
+    
+    tts_button.click(
+        process_tts,
+        inputs=translation_output,
+        outputs=audio_output
+    )
+    
+    transcribe_button.click(
+        process_transcription,
+        inputs=[audio_input, asr_language],
+        outputs=transcription_output
+    )
+
+# Launch the app
+if __name__ == "__main__":
+    demo.launch()

inference.py

@@ -0,0 +1,124 @@
+from xtransformer import Transformer
+import torch
+import torch.nn as nn
+from nepalitokenizers import SentencePiece
+from huggingface_hub import hf_hub_download
+import re
+
+# Initialize tokenizers
+tokenizer_en = SentencePiece()  # English tokenizer
+tokenizer_ne = SentencePiece()  # Nepali tokenizer
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Define special tokens and their IDs
+START_TOKEN = '<START>'
+PADDING_TOKEN = '<PADDING>'
+END_TOKEN = '<END>'
+SPECIAL_TOKENS = {
+    START_TOKEN: max(tokenizer_en.get_vocab_size(), tokenizer_ne.get_vocab_size()),
+    PADDING_TOKEN: max(tokenizer_en.get_vocab_size(), tokenizer_ne.get_vocab_size()) + 1,
+    END_TOKEN: max(tokenizer_en.get_vocab_size(), tokenizer_ne.get_vocab_size()) + 2,
+}
+
+# Update vocabulary sizes
+en_vocab_size = tokenizer_en.get_vocab_size() + len(SPECIAL_TOKENS)
+ne_vocab_size = tokenizer_ne.get_vocab_size() + len(SPECIAL_TOKENS)
+
+# Create token-to-index mappings
+english_to_index = {token: i for i, token in enumerate(tokenizer_en.get_vocab())}
+nepali_to_index = {token: i for i, token in enumerate(tokenizer_ne.get_vocab())}
+english_to_index.update(SPECIAL_TOKENS)
+nepali_to_index.update(SPECIAL_TOKENS)
+
+# Hyperparameters
+max_sequence_length = 100
+d_model = 512
+batch_size = 32
+ffn_hidden = 2048
+num_heads = 8
+drop_prob = 0.1
+encoder_layers = 6
+decoder_layers = 4
+
+# Initialize the Transformer model
+transformer = Transformer(
+    d_model, ffn_hidden, num_heads, drop_prob, encoder_layers, decoder_layers,
+    max_sequence_length, ne_vocab_size, english_to_index, nepali_to_index,
+    START_TOKEN, END_TOKEN, PADDING_TOKEN
+).to(device)
+
+# Function to encode text with special tokens
+def encode_with_special_tokens(text, tokenizer, max_sequence_length, add_start_end=True):
+    tokens = tokenizer.encode(text).ids
+    if add_start_end:
+        tokens = [SPECIAL_TOKENS[START_TOKEN]] + tokens + [SPECIAL_TOKENS[END_TOKEN]]
+    tokens = tokens[:max_sequence_length]
+    padding = [SPECIAL_TOKENS[PADDING_TOKEN]] * (max_sequence_length - len(tokens))
+    return tokens + padding
+
+# Function to decode token IDs, filtering out special tokens
+def decode_with_special_tokens(token_ids, tokenizer):
+    token_ids = [token_id for token_id in token_ids if token_id not in SPECIAL_TOKENS.values()]
+    return tokenizer.decode(token_ids)
+
+# Mask creation
+NEG_INFTY = -1e9
+def create_masks(eng_batch, decoder_input):
+    batch_size, enc_seq_length = eng_batch.size(0), eng_batch.size(1)
+    dec_seq_length = decoder_input.size(1)
+    device = eng_batch.device
+
+    encoder_padding_mask = (eng_batch == SPECIAL_TOKENS[PADDING_TOKEN]).unsqueeze(1).unsqueeze(2)
+    decoder_padding_mask_self = (decoder_input == SPECIAL_TOKENS[PADDING_TOKEN]).unsqueeze(1).unsqueeze(2)
+    look_ahead_mask = torch.triu(torch.ones(dec_seq_length, dec_seq_length, device=device), diagonal=1).bool().unsqueeze(0).unsqueeze(0)
+    decoder_padding_mask_cross = (eng_batch == SPECIAL_TOKENS[PADDING_TOKEN]).unsqueeze(1).unsqueeze(2)
+
+    encoder_mask = encoder_padding_mask * NEG_INFTY
+    decoder_self_mask = (look_ahead_mask | decoder_padding_mask_self) * NEG_INFTY
+    decoder_cross_mask = decoder_padding_mask_cross * NEG_INFTY
+
+    return encoder_mask, decoder_self_mask, decoder_cross_mask
+
+# Translation function
+def translate(sentence):
+    def is_english(text):
+        # Check if the text contains only English letters and spaces using regular expression
+        return re.match(r'^[a-zA-Z\s]+$', text) is not None
+    # Determine which model to use based on input language
+    if is_english(sentence):
+        clean_sentence = re.sub(r'[^a-zA-Z0-9\s]', '', sentence.strip()).lower() 
+        checkpoint_file = "checkpoint_en_ne.pth"
+        print('using english to nepali transformer')
+    else:
+        clean_sentence = re.sub(r'[^ऀ-ॿ\s]', '', sentence.strip())
+        checkpoint_file = "checkpoint_ne_en.pth"
+        print('using nepali to english transformer')
+    
+    # Download the checkpoint from Hugging Face Hub
+    try:
+        checkpoint_path = hf_hub_download(
+            repo_id="bishaltwr/xtransformer", 
+            filename=checkpoint_file,
+            repo_type="model"
+        )
+        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        transformer.load_state_dict(checkpoint['model_state'])
+        transformer.eval()
+    except Exception as e:
+        print(f"Error loading checkpoint: {e}")
+        return f"Translation failed: Could not load model checkpoint ({str(e)})"
+
+    with torch.no_grad():
+        eng_tokens = encode_with_special_tokens(clean_sentence, tokenizer_en, max_sequence_length)
+        eng_batch = torch.tensor([eng_tokens]).to(device)
+        ne_batch = torch.tensor([[SPECIAL_TOKENS[START_TOKEN]] + [SPECIAL_TOKENS[PADDING_TOKEN]] * (max_sequence_length - 1)]).to(device)
+
+        for i in range(1, max_sequence_length):
+            encoder_mask, decoder_mask, cross_mask = create_masks(eng_batch, ne_batch)
+            predictions = transformer(eng_batch, ne_batch, encoder_mask, decoder_mask, cross_mask)
+            next_token = torch.argmax(predictions[:, i - 1, :], dim=-1)
+            if next_token.item() == SPECIAL_TOKENS[END_TOKEN]:
+                break
+            ne_batch[0, i] = next_token
+
+    return decode_with_special_tokens(ne_batch[0].tolist(), tokenizer_ne)

requirements.txt

@@ -0,0 +1,9 @@
+gradio>=5.20.1
+torch
+transformers
+librosa
+soundfile
+gtts
+nepalitokenizers
+sounddevice
+hf_hub_download

xtransformer.py

@@ -0,0 +1,338 @@
+# import numpy as np  # Unused import
+import torch
+import math
+from torch import nn
+import torch.nn.functional as F
+from nepalitokenizers import SentencePiece
+from torch.amp import autocast  # Mixed precision
+from torch.utils.checkpoint import checkpoint  # Gradient checkpointing
+
+# Device setup
+def get_device():
+    return torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+# Efficient Scaled Dot-Product Attention
+def scaled_dot_product(q, k, v, mask=None):
+    d_k = q.size()[-1]
+    scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # Simplified attention computation
+    if mask is not None:
+        scores += mask
+    attention = F.softmax(scores, dim=-1)
+    values = torch.matmul(attention, v)
+    return values, attention
+
+# Precompute Positional Encoding
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_sequence_length):
+        super().__init__()
+        self.max_sequence_length = max_sequence_length
+        self.d_model = d_model
+        self.pe = self._create_positional_encoding()  # Precompute during initialization
+
+    def _create_positional_encoding(self):
+        position = torch.arange(self.max_sequence_length).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, self.d_model, 2).float() * (-math.log(10000.0) / self.d_model))
+        pe = torch.zeros(self.max_sequence_length, self.d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe
+
+    def forward(self, x):
+        seq_length = x.size(1)  # Handle variable sequence lengths
+        return self.pe[:seq_length, :].to(x.device)
+
+# Efficient Sentence Embedding with Caching
+class SentenceEmbedding(nn.Module):
+    def __init__(self, max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.vocab_size = len(language_to_index)
+        self.max_sequence_length = max_sequence_length
+        self.embedding = nn.Embedding(self.vocab_size, d_model)
+        self.language_to_index = language_to_index
+        self.position_encoder = PositionalEncoding(d_model, max_sequence_length)
+        self.dropout = nn.Dropout(p=0.1)
+        self.START_TOKEN = START_TOKEN
+        self.END_TOKEN = END_TOKEN
+        self.PADDING_TOKEN = PADDING_TOKEN
+        self.tokenizer = SentencePiece()
+
+class SentenceEmbedding(nn.Module):
+    def __init__(self, max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.vocab_size = len(language_to_index)
+        self.max_sequence_length = max_sequence_length
+        self.embedding = nn.Embedding(self.vocab_size, d_model)
+        self.language_to_index = language_to_index
+        self.position_encoder = PositionalEncoding(d_model, max_sequence_length)
+        self.dropout = nn.Dropout(p=0.1)
+        self.START_TOKEN = START_TOKEN
+        self.END_TOKEN = END_TOKEN
+        self.PADDING_TOKEN = PADDING_TOKEN
+        self.tokenizer = SentencePiece()
+
+    def batch_tokenize(self, batch, start_token, end_token):
+        """
+        Tokenizes a batch of sentences or processes pre-tokenized tensors.
+        
+        Args:
+            batch: A list of sentences (str) or a tensor of token IDs.
+            start_token: Whether to add a start token.
+            end_token: Whether to add an end token.
+        
+        Returns:
+            A tensor of token IDs with shape (batch_size, seq_len).
+        """
+        # If input is already a tensor, return it directly
+        if isinstance(batch, torch.Tensor):
+            return batch.to(get_device())
+
+        # Process raw text inputs
+        token_ids = []
+        for sentence in batch:
+            if not isinstance(sentence, str):
+                sentence = str(sentence).strip()
+            if not sentence:
+                sentence = self.PADDING_TOKEN
+            try:
+                tokens = self.tokenizer.encode(sentence)
+                token_ids.append(tokens.ids)
+            except Exception:
+                print(f"Error tokenizing: {sentence}")
+                token_ids.append([self.language_to_index.get(self.PADDING_TOKEN, 0)])
+
+        # Add start and end tokens if required
+        if start_token:
+            token_ids = [[self.language_to_index.get(self.START_TOKEN, self.PADDING_TOKEN)] + ids for ids in token_ids]
+        if end_token:
+            token_ids = [ids + [self.language_to_index.get(self.END_TOKEN, self.PADDING_TOKEN)] for ids in token_ids]
+
+        # Truncate sequences to max_sequence_length
+        token_ids = [ids[:self.max_sequence_length] for ids in token_ids]
+
+        # Pad sequences to max_sequence_length
+        token_ids = torch.nn.utils.rnn.pad_sequence(
+            [torch.tensor(ids, dtype=torch.long) for ids in token_ids],
+            batch_first=True,
+            padding_value=self.language_to_index.get(self.PADDING_TOKEN, 0)
+        ).to(get_device())
+
+        return token_ids
+
+    def forward(self, x, start_token, end_token):
+        """
+        Forward pass for the SentenceEmbedding module.
+        
+        Args:
+            x: Input batch (list of sentences or tensor of token IDs).
+            start_token: Whether to add a start token.
+            end_token: Whether to add an end token.
+        
+        Returns:
+            Embedded and positional-encoded output tensor.
+        """
+        # Tokenize input if it's raw text
+        if not isinstance(x, torch.Tensor):
+            x = self.batch_tokenize(x, start_token, end_token)
+
+        # Embed tokens and add positional encoding
+        x = self.embedding(x)
+        pos = self.position_encoder(x)
+        x = self.dropout(x + pos)
+        return x
+    def forward(self, x, start_token, end_token):
+    # If x is already a tensor, skip tokenization
+        if not isinstance(x, torch.Tensor):
+            x = self.batch_tokenize(x, start_token, end_token)
+        x = self.embedding(x)
+        pos = self.position_encoder(x)
+        x = self.dropout(x + pos)
+        return x
+
+# Multi-Head Attention with Efficient Matrix Operations
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.qkv_layer = nn.Linear(d_model, 3 * d_model)
+        self.linear_layer = nn.Linear(d_model, d_model)
+
+    def forward(self, x, mask):
+        batch_size, seq_length, d_model = x.size()
+        qkv = self.qkv_layer(x)
+        qkv = qkv.view(batch_size, seq_length, self.num_heads, 3 * self.head_dim)
+        qkv = qkv.permute(0, 2, 1, 3)  # (batch_size, num_heads, seq_length, 3 * head_dim)
+        q, k, v = qkv.chunk(3, dim=-1)
+        values, _ = scaled_dot_product(q, k, v, mask)  # Ignore unused variable 'attention'
+        values = values.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_length, d_model)
+        out = self.linear_layer(values)
+        return out
+
+# Multi-Head Cross Attention
+class MultiHeadCrossAttention(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        self.kv_layer = nn.Linear(d_model, 2 * d_model)
+        self.q_layer = nn.Linear(d_model, d_model)
+        self.linear_layer = nn.Linear(d_model, d_model)
+
+    def forward(self, x, y, mask):
+        batch_size, x_seq_length, _ = x.size()  # Encoder sequence length
+        batch_size, y_seq_length, _ = y.size()  # Decoder sequence length
+        
+        # Process encoder output (x) for Key/Value
+        kv = self.kv_layer(x)
+        kv = kv.view(batch_size, x_seq_length, self.num_heads, 2 * self.head_dim)
+        kv = kv.permute(0, 2, 1, 3)  # [batch, heads, x_seq, 2*head_dim]
+        k, v = kv.chunk(2, dim=-1)   # Each [batch, heads, x_seq, head_dim]
+
+        # Process decoder input (y) for Query
+        q = self.q_layer(y)
+        q = q.view(batch_size, y_seq_length, self.num_heads, self.head_dim)
+        q = q.permute(0, 2, 1, 3)    # [batch, heads, y_seq, head_dim]
+
+        # Compute attention
+        values, _ = scaled_dot_product(q, k, v, mask)
+        
+        # Reshape back to original dimensions
+        values = values.permute(0, 2, 1, 3).contiguous()
+        values = values.view(batch_size, y_seq_length, self.d_model)
+        return self.linear_layer(values)
+
+# Layer Normalization
+class LayerNormalization(nn.Module):
+    def __init__(self, parameters_shape, eps=1e-5):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(parameters_shape, eps=eps)
+
+    def forward(self, inputs):
+        return self.layer_norm(inputs)
+
+# Position-wise Feed-Forward Network
+class PositionwiseFeedForward(nn.Module):
+    def __init__(self, d_model, hidden, drop_prob=0.1):
+        super().__init__()
+        self.linear1 = nn.Linear(d_model, hidden)
+        self.linear2 = nn.Linear(hidden, d_model)
+        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(p=drop_prob)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.relu(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+# Encoder Layer with Gradient Checkpointing
+class EncoderLayer(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob):
+        super().__init__()
+        self.attention = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
+        self.norm1 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout1 = nn.Dropout(p=drop_prob)
+        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
+        self.norm2 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout2 = nn.Dropout(p=drop_prob)
+
+    def forward(self, x, self_attention_mask):
+        residual_x = x.clone()
+        x = checkpoint(self.attention, x, self_attention_mask, preserve_rng_state=True, use_reentrant=False)  # Gradient checkpointing
+        x = self.dropout1(x)
+        x = self.norm1(x + residual_x)
+        residual_x = x.clone()
+        x = checkpoint(self.ffn, x, preserve_rng_state=True, use_reentrant=False)  # Gradient checkpointing
+        x = self.dropout2(x)
+        x = self.norm2(x + residual_x)
+        return x
+
+# Sequential Encoder
+class SequentialEncoder(nn.Sequential):
+    def forward(self, *inputs):
+        x, self_attention_mask = inputs
+        for module in self._modules.values():
+            x = module(x, self_attention_mask)
+        return x
+
+# Encoder with Mixed Precision
+class Encoder(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob, encoder_layer, max_sequence_length, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.sentence_embedding = SentenceEmbedding(max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.layers = SequentialEncoder(*[EncoderLayer(d_model, ffn_hidden, num_heads, drop_prob) for _ in range(encoder_layer)])
+
+    def forward(self, x, self_attention_mask, start_token, end_token):
+        with autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):  # Mixed precision
+            x = self.sentence_embedding(x, start_token, end_token)
+            x = self.layers(x, self_attention_mask)
+        return x
+
+# Decoder Layer with Gradient Checkpointing
+class DecoderLayer(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob):
+        super().__init__()
+        self.self_attention = MultiHeadAttention(d_model=d_model, num_heads=num_heads)
+        self.layer_norm1 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout1 = nn.Dropout(p=drop_prob)
+        self.encoder_decoder_attention = MultiHeadCrossAttention(d_model=d_model, num_heads=num_heads)
+        self.layer_norm2 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout2 = nn.Dropout(p=drop_prob)
+        self.ffn = PositionwiseFeedForward(d_model=d_model, hidden=ffn_hidden, drop_prob=drop_prob)
+        self.layer_norm3 = LayerNormalization(parameters_shape=[d_model])
+        self.dropout3 = nn.Dropout(p=drop_prob)
+
+    def forward(self, x, y, self_attention_mask, cross_attention_mask):
+        _y = y.clone()
+        y = checkpoint(self.self_attention, y, self_attention_mask, preserve_rng_state=True, use_reentrant=False)  # Gradient checkpointing
+        y = self.dropout1(y)
+        y = self.layer_norm1(y + _y)
+        _y = y.clone()
+        y = checkpoint(self.encoder_decoder_attention, x, y, cross_attention_mask, preserve_rng_state=True, use_reentrant=False)  # Gradient checkpointing
+        y = self.dropout2(y)
+        y = self.layer_norm2(y + _y)
+        _y = y.clone()
+        y = checkpoint(self.ffn, y, preserve_rng_state=True, use_reentrant=False)  # Gradient checkpointing
+        y = self.dropout3(y)
+        y = self.layer_norm3(y + _y)
+        return y
+
+# Sequential Decoder
+class SequentialDecoder(nn.Sequential):
+    def forward(self, *inputs):
+        x, y, self_attention_mask, cross_attention_mask = inputs
+        for module in self._modules.values():
+            y = module(x, y, self_attention_mask, cross_attention_mask)
+        return y
+
+# Decoder with Mixed Precision
+class Decoder(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob, decoder_layer, max_sequence_length, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.sentence_embedding = SentenceEmbedding(max_sequence_length, d_model, language_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.layers = SequentialDecoder(*[DecoderLayer(d_model, ffn_hidden, num_heads, drop_prob) for _ in range(decoder_layer)])
+
+    def forward(self, x, y, self_attention_mask, cross_attention_mask, start_token, end_token):
+        with autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):  # Mixed precision
+            y = self.sentence_embedding(y, start_token, end_token)
+            y = self.layers(x, y, self_attention_mask, cross_attention_mask)
+        return y
+
+# Transformer with Mixed Precision and Gradient Checkpointing
+class Transformer(nn.Module):
+    def __init__(self, d_model, ffn_hidden, num_heads, drop_prob, encoder_layer, decoder_layer, max_sequence_length, ne_vocab_size, english_to_index, nepali_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN):
+        super().__init__()
+        self.encoder = Encoder(d_model, ffn_hidden, num_heads, drop_prob, encoder_layer, max_sequence_length, english_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.decoder = Decoder(d_model, ffn_hidden, num_heads, drop_prob, decoder_layer, max_sequence_length, nepali_to_index, START_TOKEN, END_TOKEN, PADDING_TOKEN)
+        self.linear = nn.Linear(d_model, ne_vocab_size)
+        self.device = get_device()
+
+    def forward(self, x, y, encoder_self_attention_mask=None, decoder_self_attention_mask=None, decoder_cross_attention_mask=None, enc_start_token=False, enc_end_token=False, dec_start_token=False, dec_end_token=False):
+        with autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):  # Mixed precision
+            x = self.encoder(x, encoder_self_attention_mask, enc_start_token, enc_end_token)
+            out = self.decoder(x, y, decoder_self_attention_mask, decoder_cross_attention_mask, dec_start_token, dec_end_token)
+            out = self.linear(out)
+        return out