Add gradio app file

app.py

@@ -0,0 +1,262 @@
+import math
+
+import gradio as gr
+import lightning as L
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+from tokenizers import Tokenizer
+
+
+class Translator:
+    def __init__(
+        self,
+        src_tokenizer_ckpt_path,
+        tgt_tokenizer_ckpt_path,
+        model_ckpt_path,
+    ):
+        self.src_tokenizer = Tokenizer.from_file(src_tokenizer_ckpt_path)
+        self.tgt_tokenizer = Tokenizer.from_file(tgt_tokenizer_ckpt_path)
+
+        self.src_tokenizer.model.dropout = 0
+        self.tgt_tokenizer.model.dropout = 0
+
+        self.model = TransformerSeq2Seq.load_from_checkpoint(
+            model_ckpt_path,
+            map_location="cpu",
+        )
+        self.model.eval()
+
+    def predict(self, src):
+        tokenized_text = self.src_tokenizer.encode(src)
+        src = torch.LongTensor(tokenized_text.ids).view(-1, 1)
+        tgt = self.model.greedy_decode(src, max_len=100)
+        tgt = tgt.squeeze(1).tolist()
+        tgt_text = self.tgt_tokenizer.decode(tgt)
+        return tgt_text
+
+
+def generate_square_subsequent_mask(sz):
+    mask = (torch.triu(torch.ones((sz, sz))) == 1).transpose(0, 1)
+    mask = (
+        mask.float()
+        .masked_fill(mask == 0, float("-inf"))
+        .masked_fill(mask == 1, float(0.0))
+    )
+    return mask
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, embedding_dim, dropout, maxlen=5000):
+        super(PositionalEncoding, self).__init__()
+        den = torch.exp(
+            -torch.arange(0, embedding_dim, 2) * math.log(10000) / embedding_dim
+        )
+        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
+        pos_embedding = torch.zeros((maxlen, embedding_dim))
+        pos_embedding[:, 0::2] = torch.sin(pos * den)
+        pos_embedding[:, 1::2] = torch.cos(pos * den)
+        pos_embedding = pos_embedding.unsqueeze(-2)
+
+        self.dropout = nn.Dropout(dropout)
+        self.register_buffer("pos_embedding", pos_embedding)
+
+    def forward(self, token_embedding):
+        return self.dropout(
+            token_embedding + self.pos_embedding[: token_embedding.size(0), :]
+        )
+
+
+class TransformerSeq2Seq(L.LightningModule):
+    def __init__(
+        self,
+        src_vocab_size,
+        tgt_vocab_size,
+        embedding_dim=512,
+        hidden_dim=512,
+        dropout=0.1,
+        nhead=8,
+        num_layers=3,
+        batch_size=32,
+        lr=1e-4,
+        weight_decay=1e-4,
+        sos_idx=1,
+        eos_idx=2,
+        padding_idx=3,
+    ):
+        super().__init__()
+        self.save_hyperparameters()
+
+        self.src_embedding = nn.Embedding(
+            src_vocab_size,
+            embedding_dim,
+            padding_idx=padding_idx,
+        )
+        self.tgt_embedding = nn.Embedding(
+            tgt_vocab_size,
+            embedding_dim,
+            padding_idx=padding_idx,
+        )
+        self.positional_encoding = PositionalEncoding(
+            embedding_dim=embedding_dim,
+            dropout=dropout,
+        )
+        self.transformer = nn.Transformer(
+            d_model=embedding_dim,
+            nhead=nhead,
+            num_encoder_layers=num_layers,
+            num_decoder_layers=num_layers,
+            dim_feedforward=hidden_dim,
+            dropout=dropout,
+        )
+        self.fc = nn.Linear(embedding_dim, tgt_vocab_size)
+
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+        self.criteria = nn.CrossEntropyLoss()
+
+    def forward(
+        self,
+        src,
+        tgt,
+        src_mask,
+        tgt_mask,
+        src_padding_mask,
+        tgt_padding_mask,
+    ):
+        src = self.src_embedding(src) * (self.hparams.embedding_dim**0.5)
+        tgt = self.tgt_embedding(tgt) * (self.hparams.embedding_dim**0.5)
+        src = self.positional_encoding(src)
+        tgt = self.positional_encoding(tgt)
+        out = self.transformer(
+            src,
+            tgt,
+            src_mask=src_mask,
+            tgt_mask=tgt_mask,
+            src_key_padding_mask=src_padding_mask,
+            tgt_key_padding_mask=tgt_padding_mask,
+        )
+        out = self.fc(out)
+        return out
+
+    def greedy_decode(self, src, max_len):
+        src = self.src_embedding(src) * (self.hparams.embedding_dim**0.5)
+        src = self.positional_encoding(src)
+        memory = self.transformer.encoder(src)
+        ys = torch.ones(1, 1).fill_(self.hparams.sos_idx).type(torch.long)
+        for i in range(max_len - 1):
+            tgt = self.tgt_embedding(ys) * (self.hparams.embedding_dim**0.5)
+            tgt = self.positional_encoding(tgt)
+            tgt_mask = generate_square_subsequent_mask(ys.size(0)).type(torch.bool)
+            out = self.transformer.decoder(
+                tgt,
+                memory,
+                tgt_mask=tgt_mask,
+            )
+            out = self.fc(out)
+            out = out.transpose(0, 1)[:, -1]
+            prob = out.softmax(dim=-1)
+            _, next_word = torch.max(prob, dim=1)
+            next_word = next_word.item()
+            ys = torch.cat(
+                [ys, torch.ones(1, 1).fill_(next_word).type(torch.long)],
+                dim=0,
+            )
+
+            if next_word == self.hparams.eos_idx:
+                break
+
+        return ys
+
+    def training_step(self, batch, batch_idx):
+        src, tgt, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = batch
+        tgt_input = tgt[:-1, :]
+        logits = self(
+            src,
+            tgt_input,
+            src_mask,
+            tgt_mask,
+            src_padding_mask,
+            tgt_padding_mask,
+        )
+        tgt_out = tgt[1:, :]
+        loss = self.criteria(
+            logits.reshape(-1, logits.shape[-1]),
+            tgt_out.reshape(-1),
+        )
+        self.log("train_loss", loss, batch_size=self.hparams.batch_size)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        src, tgt, src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = batch
+        tgt_input = tgt[:-1, :]
+        logits = self(
+            src,
+            tgt_input,
+            src_mask,
+            tgt_mask,
+            src_padding_mask,
+            tgt_padding_mask,
+        )
+        tgt_out = tgt[1:, :]
+        loss = self.criteria(
+            logits.reshape(-1, logits.shape[-1]),
+            tgt_out.reshape(-1),
+        )
+        self.log("val_loss", loss, batch_size=self.hparams.batch_size)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(
+            self.parameters(),
+            lr=self.hparams.lr,
+            weight_decay=self.hparams.weight_decay,
+        )
+        return {
+            "optimizer": optimizer,
+            "lr_scheduler": {
+                "scheduler": torch.optim.lr_scheduler.OneCycleLR(
+                    optimizer=optimizer,
+                    max_lr=self.hparams.lr,
+                    total_steps=self.trainer.estimated_stepping_batches,
+                ),
+                "interval": "step",
+            },
+        }
+
+
+src_tokenizer_ckpt_path = hf_hub_download(
+    repo_id="SatwikKambham/opus100-en-hi-transformer",
+    filename="tokenizer-en.json",
+)
+tgt_tokenizer_ckpt_path = hf_hub_download(
+    repo_id="SatwikKambham/opus100-en-hi-transformer",
+    filename="tokenizer-hi.json",
+)
+model_ckpt_path = hf_hub_download(
+    repo_id="SatwikKambham/opus100-en-hi-transformer",
+    filename="transformer.ckpt",
+)
+classifier = Translator(
+    src_tokenizer_ckpt_path,
+    tgt_tokenizer_ckpt_path,
+    model_ckpt_path,
+)
+interface = gr.Interface(
+    fn=classifier.predict,
+    inputs=gr.components.Textbox(
+        label="Source Language (English)",
+        placeholder="Enter text here...",
+    ),
+    outputs=gr.components.Textbox(
+        label="Target Language (Hindi)",
+        placeholder="Translation",
+    ),
+    examples=[
+        ["Hi how are you?"],
+        ["Today is a very important day."],
+        ["I like playing the guitar."],
+    ],
+)
+interface.launch()