import gradio as gr
from transformers import pipeline
import torch

generator = pipeline('text-generation', model='Sigurdur/icebreaker')

def generate(text):
    result = generator(text, max_length=30, num_return_sequences=1)
    return result[0]["generated_text"]

examples = ['<Afmæli>', '<Afríka>', '<Afþreying>', '<Almannavarnir>', '<Alþingiskosningar>', '<Asía>', '<Atvinnulíf>', '<Atvinnumál>', '<Auðlesið efni>', '<Aðsent>', '<Bandaríkin>', '<Barnaefni>', '<Bretland>', '<Bæjar- og sveitarfélög>', '<Bílar>', '<Bíó og sjónvarp>', '<Bókmenntir>', '<Covid>', '<Covid-19>', '<Dóms- og lögreglumál>', '<Dýr>', '<Esb>', '<Efnahagsmál>', '<Erlendar fréttir>', '<Eurovision>', '<Evrópa>', '<Eyjaálfa>', '<Fasteignir og húsæðsmál>', '<Ferðalög>', '<Ferðamál>', '<Fjölskyldan>', '<Flug>', '<Flóttafólk og innflytjendamál>', '<Fréttaskýringar>', '<Fréttir>', '<Fólk>', '<Fótbolti>', '<Fótbolti - ísland>', '<Gagnrýni>', '<Garðrækt>', '<Golf>', '<Handbolti>', '<Handbolti - ísland>', '<Hannyrðir>', '<Heilbrigðismál>', '<Heilsa og útlit>', '<Heimilið>', '<Hryðjuverk>', '<Innlent>', '<Jafnréttismál>', '<Jól>', '<Kjaramál>', '<Klaustursmálið>', '<Kosningar>', '<Kvikmyndir>', '<Kynning>', '<Körfubolti>', '<Körfubolti - ísland>', '<Landbúnaður>', '<Landshlutar>', '<Leiklist>', '<Leiðarar>', '<Lesbók morgunblaðsins>', '<Loftslagsmál>', '<Lífsstíll>', '<Mannlíf>', '<Matur og drykkur>', '<Menning>', '<Menntamál>', '<Minningargreinar>', '<Mið- og suður-ameríka>', '<Mið-austurlönd>', '<Myndlist>', '<Neytendamál>', '<Norður-ameríka>', '<Norðurlönd>', '<Náttúra og umhverfi>', '<Samgöngumál>', '<Sjávarútvegsmál>', '<Sjónvarp>', '<Skopmyndir>', '<Slys og eldsvoðar>', '<Smartland>', '<Stjórnmál>', '<Stríð>', '<Tilkynningar>', '<Trúarbrögð>', '<Tækni og vísindi>', '<Tíska>', '<Tónlist>', '<Tölvuleikir>', '<Umfjöllun>', '<Umræða>', '<Uppljóstrun>', '<Veiði>', '<Veður>', '<Veður og náttúruhamfarir>', '<Viðhorf>', '<Viðskipti>', '<Viðtöl>', '<Í dag - mbl.is>', '<Íslensk stjórnmál>', '<Íslenskt mál>', '<Íþróttir>', '<Óflokkað>', '<Ólympíuleikar>']

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(label="Generated Text"),
    examples=examples,
    title="Generate news articles in Icelandic with Icebreaker",
    description="Author: Sigurdur Haukur Birgisson",

)

demo.launch()