File size: 1,498 Bytes
ddf9cdc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import gradio as gr
# Check if CUDA (GPU) is available
import torch
from transformers import T5ForConditionalGeneration, PreTrainedTokenizerFast
# Define the path to the checkpoint directory
checkpoint_dir = "/home/only_sainaa/home/Huggingface/marianmt_conversion/results/checkpoint-221496"
# Load the model
model = T5ForConditionalGeneration.from_pretrained(checkpoint_dir)
model.eval()
# Load the tokenizer using PreTrainedTokenizerFast
tokenizer = PreTrainedTokenizerFast.from_pretrained(checkpoint_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Move the model to the same device (GPU or CPU)
model.to(device)
# Function to perform translation using the model
def translate_text(input_text):
# Tokenize the input text
inputs = tokenizer(input_text, return_tensors="pt")
# Move the input tensors to the same device as the model
inputs = {k: v.to(device) for k, v in inputs.items() if k in ['input_ids', 'attention_mask']}
# Generate translation
outputs = model.generate(**inputs)
# Decode the output to human-readable text
translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return translated_text
# Create a Gradio interface
gr_interface = gr.Interface(
fn=translate_text,
inputs="text",
outputs="text",
title="Mongolian Cyrillic to Mongolian Script Model",
description="Enter text in Mongolian Cyrillic"
)
# Launch the Gradio interface
gr_interface.launch() |