Upload 4 files

DockerFile

@@ -0,0 +1,10 @@
+FROM python:3.9-slim
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+COPY ./app /code/app
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

app/main.py

@@ -0,0 +1,102 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+import torch
+import torch.nn.functional as F
+from tokenizers import Tokenizer
+from huggingface_hub import hf_hub_download
+
+from .model_def import BuildTransformer
+
+
+app = FastAPI(title="Hindi-English Translator API")
+model = None
+tokenizer = None
+device = torch.device("cpu") 
+
+
+class TranslationRequest(BaseModel):
+    text: str
+
+class TranslationResponse(BaseModel):
+    translated_text: str
+
+@app.on_event("startup")
+def load_assets():
+
+    global model, tokenizer, device
+
+    model_file = hf_hub_download(repo_id="Kush26/Transformer_Translation", filename="model.pth")
+    tokenizer_file = hf_hub_download(repo_id="Kush26/Transformer_Translation", filename="hindi-english_bpe_tokenizer.json")
+    
+    tokenizer = Tokenizer.from_file(tokenizer_file)
+    vocab_size = tokenizer.get_vocab_size()
+
+    config = {
+        "d_model": 256,
+        "num_layers": 6,
+        "num_heads": 8,
+        "d_ff": 2048,
+        "dropout": 0.1,
+        "max_seq_len": 512,
+    }
+
+    model = BuildTransformer(
+        src_vocab_size=vocab_size,
+        trg_vocab_size=vocab_size,
+        src_seq_len=config["max_seq_len"],
+        trg_seq_len=config["max_seq_len"],
+        d_model=config["d_model"],
+        N=config["num_layers"],
+        h=config["num_heads"],
+        dropout=config["dropout"],
+        d_ff=config["d_ff"]
+    ).to(device)
+
+    # 5. Load the trained weights
+    checkpoint = torch.load(model_file, map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval() # Set model to evaluation mode
+
+    print("✅ Model and Tokenizer loaded successfully!")
+
+
+def greedy_decode(sentence: str, max_len=100):
+    PAD_token = tokenizer.token_to_id('[PAD]')
+    
+    model.eval()
+    src_ids = [tokenizer.token_to_id('[SOS]')] + tokenizer.encode(sentence).ids + [tokenizer.token_to_id('[EOS]')]
+    src_tensor = torch.tensor(src_ids).unsqueeze(0).to(device)
+    src_mask = (src_tensor != PAD_token).unsqueeze(1).unsqueeze(2)
+
+    with torch.no_grad():
+        encoder_output = model.encode(src_tensor, src_mask)
+
+    tgt_tokens = [tokenizer.token_to_id('[SOS]')]
+    
+    for _ in range(max_len):
+        tgt_tensor = torch.tensor(tgt_tokens).unsqueeze(0).to(device)
+        trg_mask_padding = (tgt_tensor != PAD_token).unsqueeze(1).unsqueeze(2)
+        subsequent_mask = torch.tril(torch.ones(1, tgt_tensor.size(1), tgt_tensor.size(1), device=device)).bool()
+        trg_mask = trg_mask_padding & subsequent_mask
+
+        with torch.no_grad():
+            decoder_output = model.decode(encoder_output, src_mask, tgt_tensor, trg_mask)
+            logits = model.project(decoder_output)
+        
+        pred_token = logits.argmax(dim=-1)[0, -1].item()
+        tgt_tokens.append(pred_token)
+        
+        if pred_token == tokenizer.token_to_id('[EOS]'):
+            break
+            
+    return tokenizer.decode(tgt_tokens, skip_special_tokens=True)
+
+@app.get("/")
+def read_root():
+    return {"message": "Welcome to the Hindi-English Translator API"}
+
+@app.post("/translate/greedy", response_model=TranslationResponse)
+def translate_greedy_endpoint(request: TranslationRequest):
+
+    translated_text = greedy_decode(request.text)
+    return {"translated_text": translated_text}

app/model_def.py

@@ -0,0 +1,243 @@
+import torch
+import torch.nn as nn
+import math
+
+class InputEmbedding(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.d_model = d_model
+        self.vocab_size = vocab_size
+        self.embed = nn.Embedding(vocab_size, d_model)
+
+    def forward(self, x):
+        return self.embed(x) * math.sqrt(self.d_model)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, seq_len: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.seq_len = seq_len
+        self.dropout = nn.Dropout(dropout)
+
+        # Create a matrix of shape (seq_len, d_model)
+        pe = torch.zeros(seq_len, d_model)
+        
+        # Create a vector of shape (seq_len, 1)
+        position = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        
+        # Apply sin to even indices
+        pe[:, 0::2] = torch.sin(position * div_term)
+        # Apply cos to odd indices
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        # Add a batch dimension
+        pe = pe.unsqueeze(0) # (1, seq_len, d_model)
+
+        # Register 'pe' as a buffer, so it's not a model parameter
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        x = x + (self.pe[:, :x.shape[1], :]).requires_grad_(False)
+        return self.dropout(x)
+
+class LayerNorm(nn.Module):
+    def __init__(self, d_model: int, epsilon: float = 1e-6):
+        super().__init__()
+        self.epsilon = epsilon
+        self.gamma = nn.Parameter(torch.ones(d_model)) # Multiplicative
+        self.beta = nn.Parameter(torch.zeros(d_model))  # Additive
+
+    def forward(self, x):
+        mean = x.mean(dim=-1, keepdim=True)
+        std = x.std(dim=-1, keepdim=True)
+        return self.gamma * (x - mean) / (std + self.epsilon) + self.beta
+
+class FeedForward(nn.Module):
+    def __init__(self, d_model: int, d_ff: int, dropout: float):
+        super().__init__()
+        self.layer1 = nn.Linear(d_model, d_ff)
+        self.layer2 = nn.Linear(d_ff, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, d_ff) -> (Batch, Seq_Len, d_model)
+        return self.layer2(self.dropout(torch.relu(self.layer1(x))))
+
+class MHA(nn.Module):
+    def __init__(self, d_model: int, h: int, dropout: float):
+        super().__init__()
+        self.d_model = d_model
+        self.h = h
+        assert d_model % h == 0, "d_model must be divisible by h"
+
+        self.d_k = d_model // h
+        self.w_q = nn.Linear(d_model, d_model)
+        self.w_k = nn.Linear(d_model, d_model)
+        self.w_v = nn.Linear(d_model, d_model)
+        self.w_o = nn.Linear(d_model, d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    @staticmethod
+    def attention(query, key, value, mask, dropout: nn.Dropout):
+        d_k = query.shape[-1]
+        attention_scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
+        if mask is not None:
+            attention_scores = attention_scores.masked_fill(mask == 0, -1e9)
+        
+        attention_scores = attention_scores.softmax(dim=-1)
+        if dropout is not None:
+            attention_scores = dropout(attention_scores)
+        
+        return (attention_scores @ value), attention_scores
+
+    def forward(self, q, k, v, mask):
+        query = self.w_q(q) # (Batch, Seq_Len, d_model)
+        key = self.w_k(k)   # (Batch, Seq_Len, d_model)
+        value = self.w_v(v) # (Batch, Seq_Len, d_model)
+
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, h, d_k) -> (Batch, h, Seq_Len, d_k)
+        query = query.view(query.shape[0], query.shape[1], self.h, self.d_k).transpose(1, 2)
+        key = key.view(key.shape[0], key.shape[1], self.h, self.d_k).transpose(1, 2)
+        value = value.view(value.shape[0], value.shape[1], self.h, self.d_k).transpose(1, 2)
+
+        x, self.attention_scores = MHA.attention(query, key, value, mask, self.dropout)
+
+        # (Batch, h, Seq_Len, d_k) -> (Batch, Seq_Len, h, d_k) -> (Batch, Seq_Len, d_model)
+        x = x.transpose(1, 2).contiguous().view(x.shape[0], -1, self.h * self.d_k)
+
+        return self.w_o(x)
+
+class SkipConnection(nn.Module):
+    def __init__(self, d_model: int, dropout: float):
+        super().__init__()
+        self.dropout = nn.Dropout(dropout)
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, sublayer):
+        # Pre-Norm architecture
+        return x + self.dropout(sublayer(self.norm(x)))
+
+class EncoderBlock(nn.Module):
+    def __init__(self, self_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
+        super().__init__()
+        self.self_attention = self_attention
+        self.ffn = ffn
+        self.skip_connections = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(2)])
+
+    def forward(self, x, src_mask):
+        x = self.skip_connections[0](x, lambda x: self.self_attention(x, x, x, src_mask))
+        x = self.skip_connections[1](x, self.ffn)
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, d_model: int, layers: nn.ModuleList):
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, mask):
+        for layer in self.layers:
+            x = layer(x, mask)
+        return self.norm(x)
+
+class DecoderBlock(nn.Module):
+    def __init__(self, self_attention: MHA, cross_attention: MHA, ffn: FeedForward, d_model: int, dropout: float):
+        super().__init__()
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.ffn = ffn
+        self.skip_connections = nn.ModuleList([SkipConnection(d_model, dropout) for _ in range(3)])
+
+    def forward(self, x, encoder_output, src_mask, trg_mask):
+        x = self.skip_connections[0](x, lambda x: self.self_attention(x, x, x, trg_mask))
+        x = self.skip_connections[1](x, lambda x: self.cross_attention(x, encoder_output, encoder_output, src_mask))
+        x = self.skip_connections[2](x, self.ffn)
+        return x
+
+class Decoder(nn.Module):
+    def __init__(self, d_model: int, layers: nn.ModuleList):
+        super().__init__()
+        self.layers = layers
+        self.norm = LayerNorm(d_model)
+
+    def forward(self, x, encoder_output, src_mask, trg_mask):
+        for layer in self.layers:
+            x = layer(x, encoder_output, src_mask, trg_mask)
+        return self.norm(x)
+
+class Output(nn.Module):
+    def __init__(self, d_model: int, vocab_size: int):
+        super().__init__()
+        self.proj = nn.Linear(d_model, vocab_size)
+
+    def forward(self, x):
+        # (Batch, Seq_Len, d_model) -> (Batch, Seq_Len, vocab_size)
+        return self.proj(x)
+
+class Transformer(nn.Module):
+    def __init__(self, encoder: Encoder, decoder: Decoder, src_embed: InputEmbedding, trg_embed: InputEmbedding, src_pos: PositionalEncoding, trg_pos: PositionalEncoding, output: Output):
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.src_embed = src_embed
+        self.trg_embed = trg_embed
+        self.src_pos = src_pos
+        self.trg_pos = trg_pos
+        self.output_layer = output
+
+    def encode(self, src, src_mask):
+        src = self.src_embed(src)
+        src = self.src_pos(src)
+        return self.encoder(src, src_mask)
+
+    def decode(self, encoder_output, src_mask, trg, trg_mask):
+        trg = self.trg_embed(trg)
+        trg = self.trg_pos(trg)
+        return self.decoder(trg, encoder_output, src_mask, trg_mask)
+
+    def project(self, x):
+        return self.output_layer(x)
+
+def BuildTransformer(src_vocab_size: int, trg_vocab_size: int, src_seq_len: int, trg_seq_len: int, d_model: int = 512, N: int = 6, h: int = 8, dropout: float = 0.1, d_ff: int = 2048) -> Transformer:
+    # Create the embedding layers
+    src_embed = InputEmbedding(d_model, src_vocab_size)
+    trg_embed = InputEmbedding(d_model, trg_vocab_size)
+
+    # Create the positional encoding layers
+    src_pos = PositionalEncoding(d_model, src_seq_len, dropout)
+    trg_pos = PositionalEncoding(d_model, trg_seq_len, dropout)
+
+    # Create the encoder blocks
+    encoder_blocks = []
+    for _ in range(N):
+        encoder_self_attention = MHA(d_model, h, dropout)
+        ffn = FeedForward(d_model, d_ff, dropout)
+        encoder_block = EncoderBlock(encoder_self_attention, ffn, d_model, dropout)
+        encoder_blocks.append(encoder_block)
+
+    # Create the decoder blocks
+    decoder_blocks = []
+    for _ in range(N):
+        decoder_self_attention = MHA(d_model, h, dropout)
+        cross_attention = MHA(d_model, h, dropout)
+        ffn = FeedForward(d_model, d_ff, dropout)
+        decoder_block = DecoderBlock(decoder_self_attention, cross_attention, ffn, d_model, dropout)
+        decoder_blocks.append(decoder_block)
+    
+    # Create the encoder and decoder
+    encoder = Encoder(d_model, nn.ModuleList(encoder_blocks))
+    decoder = Decoder(d_model, nn.ModuleList(decoder_blocks))
+
+    # Create the projection layer
+    projection = Output(d_model, trg_vocab_size)
+
+    # Create the transformer
+    transformer = Transformer(encoder, decoder, src_embed, trg_embed, src_pos, trg_pos, projection)
+
+    # Initialize parameters
+    for p in transformer.parameters():
+        if p.dim() > 1:
+            nn.init.xavier_uniform_(p)
+    
+    return transformer

requirements.txt

@@ -0,0 +1,5 @@
+fastapi
+uvicorn[standard]
+torch
+tokenizers
+huggingface_hub