import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelWithLMHead

tokenizer = AutoTokenizer.from_pretrained("mmcquade11/autonlp-reuters-summarization-34018133")
model = AutoModelWithLMHead.from_pretrained("mmcquade11/autonlp-reuters-summarization-34018133")

def summarize(text):
    input_ids = torch.tensor(tokenizer.encode(text, add_special_tokens=True)).unsqueeze(0)
    summary_ids = model.generate(input_ids, num_beams=4, max_length=100, early_stopping=True)
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)

def summarize_text(text):
    return summarize(text)

iface = gr.Interface(summarize_text, "textbox", "label")
if __name__ == "__main__":
    iface.launch()