File size: 1,010 Bytes
1a8c724 4d3d295 1a8c724 4d3d295 1a8c724 e9e5a55 |
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 |
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("./models/checkpoint-15000/")
def text_processing(text):
inputs = [text]
# Tokenize and prepare the inputs for model
input_ids = tokenizer(inputs, return_tensors="pt", max_length=512, truncation=True, padding="max_length").input_ids.to(device)
attention_mask = tokenizer(inputs, return_tensors="pt", max_length=512, truncation=True, padding="max_length").attention_mask.to(device)
# Generate prediction
output = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=512)
# Decode the prediction
decoded_output = [tokenizer.decode(ids, skip_special_tokens=True) for ids in output]
return decoded_output[0]
iface = gr.Interface(fn = text_processing, inputs='text', outputs=['text'], title='test', description='test space')
iface.launch(inline=False) |