Spaces:
Sleeping
Sleeping
| 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() | |