|
import gradio as gr |
|
import torch |
|
from transformers import BertTokenizerFast, EncoderDecoderModel |
|
|
|
device = 'cuda' if torch.cuda.is_available() else 'cpu' |
|
ckpt = 'mrm8488/bert2bert_shared-spanish-finetuned-summarization' |
|
tokenizer = BertTokenizerFast.from_pretrained(ckpt) |
|
model = EncoderDecoderModel.from_pretrained(ckpt).to(device) |
|
|
|
def generate_summary(text): |
|
inputs = tokenizer([text], padding="max_length", truncation=True, max_length=512, return_tensors="pt") |
|
input_ids = inputs.input_ids.to(device) |
|
attention_mask = inputs.attention_mask.to(device) |
|
output = model.generate(input_ids, attention_mask=attention_mask) |
|
return tokenizer.decode(output[0], skip_special_tokens=True) |
|
|
|
|
|
textbox = gr.Textbox(placeholder="Introduzca texto a resumir ...", lines=4) |
|
demo = gr.Interface(fn=generate_summary, inputs=textbox,outputs="text") |
|
|
|
demo.launch() |
|
|