GermanToEnglish / translator_app.py
neelimapreeti297's picture
Upload translator_app.py
3b60890 verified
raw
history blame
1.68 kB
import gradio as gr
import torch
from torchtext.data.utils import get_tokenizer
from torchtext.vocab import build_vocab_from_iterator
from torchtext.datasets import Multi30k
from torch import Tensor
from typing import Iterable, List
# Define your model, tokenizer, and other necessary components here
# Ensure you have imported all necessary libraries
# Load your transformer model
model = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load('./transformer_model.pth', map_location=device))
model.eval()
def translate(model: torch.nn.Module, src_sentence: str):
model.eval()
src = text_transform[SRC_LANGUAGE](src_sentence).view(-1, 1)
num_tokens = src.shape[0]
src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
tgt_tokens = greedy_decode(
model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
return " ".join(vocab_transform[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("<bos>", "").replace("<eos>", "")
if __name__ == "__main__":
# Create the Gradio interface
iface = gr.Interface(
fn=translate, # Specify the translation function as the main function
inputs=[
gr.inputs.Textbox(label="Text")
],
outputs=["text"], # Define the output type as text
cache_examples=False, # Disable caching of examples
title="germanToenglish", # Set the title of the interface
)
# Launch the interface
iface.launch(share=True)