Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,204 Bytes
60af901 663f046 8b62553 4298bfe 8b62553 4298bfe 8b62553 4298bfe 663f046 4298bfe 663f046 4298bfe 60af901 8b62553 4298bfe 8b62553 60af901 8b62553 4298bfe 8b62553 60af901 8b62553 4298bfe 8b62553 4298bfe 8b62553 bdb369b 23bdd0d 8b62553 ffb5497 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import spaces
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import spacy
class ModelSingleton:
_instance = None
def __new__(cls, *args, **kwargs):
if not cls._instance:
cls._instance = super(ModelSingleton, cls).__new__(cls, *args, **kwargs)
return cls._instance
def __init__(self):
if not hasattr(self, 'initialized'):
self.nlp_en = spacy.load("en_core_web_sm")
self.nlp_it = spacy.load("it_core_news_sm")
# Load translation models and tokenizers
self.tokenizer_en_it = AutoTokenizer.from_pretrained("LeonardPuettmann/Quadrifoglio-mt-en-it")
self.model_en_it = AutoModelForSeq2SeqLM.from_pretrained("LeonardPuettmann/Quadrifoglio-mt-en-it", torch_dtype=torch.bfloat16)
self.tokenizer_it_en = AutoTokenizer.from_pretrained("LeonardPuettmann/Quadrifoglio-mt-it-en")
self.model_it_en = AutoModelForSeq2SeqLM.from_pretrained("LeonardPuettmann/Quadrifoglio-mt-it-en", torch_dtype=torch.bfloat16)
self.initialized = True
model_singleton = ModelSingleton()
@spaces.GPU(duration=30)
def generate_response_en_it(input_text):
input_ids = model_singleton.tokenizer_en_it("translate English to Italian: " + input_text, return_tensors="pt").input_ids
output = model_singleton.model_en_it.generate(input_ids, max_new_tokens=256)
return model_singleton.tokenizer_en_it.decode(output[0], skip_special_tokens=True)
@spaces.GPU(duration=30)
def generate_response_it_en(input_text):
input_ids = model_singleton.tokenizer_it_en("translate Italian to English: " + input_text, return_tensors="pt").input_ids
output = model_singleton.model_it_en.generate(input_ids, max_new_tokens=256)
return model_singleton.tokenizer_it_en.decode(output[0], skip_special_tokens=True)
@spaces.GPU(duration=30)
def translate_text(input_text, direction):
if direction == "en-it":
nlp = model_singleton.nlp_en
generate_response = generate_response_en_it
elif direction == "it-en":
nlp = model_singleton.nlp_it
generate_response = generate_response_it_en
else:
return "Invalid direction selected."
doc = nlp(input_text)
sentences = [sent.text for sent in doc.sents]
sentence_translations = []
for sentence in sentences:
sentence_translation = generate_response(sentence)
sentence_translations.append(sentence_translation)
full_translation = " ".join(sentence_translations)
return full_translation
# Create the Gradio interface
iface = gr.Interface(
fn=translate_text,
inputs=[gr.Textbox(lines=5, placeholder="Enter text to translate...", label="Input Text"),
gr.Dropdown(choices=["en-it", "it-en"], label="Translation Direction")],
outputs=gr.Textbox(lines=5, label="Translation"),
description="This space is running on ZERO GPU. Initilization might take a couple of seconds the first time. This spaces uses the Quadrifoglio models for it-en and en-it text translation tasks."
)
# Launch the interface
iface.launch()
|