Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import torch | |
| from transformers import RobertaTokenizerFast, BertTokenizerFast, EncoderDecoderModel | |
| LANGUAGES = ["fr", "de", "tu", "es"] | |
| models = dict() | |
| tokenizers = dict() | |
| models_paths = dict() | |
| models_paths["fr"] = "mrm8488/camembert2camembert_shared-finetuned-french-summarization" | |
| models_paths["de"] = "mrm8488/bert2bert_shared-german-finetuned-summarization" | |
| models_paths["tu"] = "mrm8488/bert2bert_shared-turkish-summarization" | |
| models_paths["es"] = "Narrativa/bsc_roberta2roberta_shared-spanish-finetuned-mlsum-summarization" | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| for lang in LANGUAGES: | |
| tokenizers[lang] = RobertaTokenizerFast.from_pretrained(models_paths[lang]) if lang in ["fr", "es"] else BertTokenizerFast.from_pretrained(models_paths[lang]) | |
| models[lang] = EncoderDecoderModel.from_pretrained(models_paths[lang]).to(device) | |
| def summarize(lang, text): | |
| tokenizer = tokenizers[lang] | |
| model = models[lang] | |
| 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) | |
| theme = "darkgrass" | |
| title = "Multilingual Summarization model (MLSUM)" | |
| description = "Gradio Demo for Summarization models trained on MLSUM dataset by Manuel Romero" | |
| article = "<p style='text-align: center'><a href='https://hf.co/mrm8488' target='_blank'>More models</a></p>" | |
| gr.Interface(fn=summarize, inputs=[gr.inputs.Radio(LANGUAGES), gr.inputs.Textbox( | |
| lines=7, label="Input Text")], outputs="text", theme=theme, title=title, description=description, article=article, enable_queue=True).launch(inline=False) | |