import streamlit as st
from transformers import AutoTokenizer, AutoModelWithLMHead, GPT2Tokenizer, GPT2Model, FlaxGPT2LMHeadModel, GPT2LMHeadModel, pipeline, set_seed
import torch

#===========================================#
#        Loads Model and Pipeline        #
#===========================================#

# tokenizer = AutoTokenizer.from_pretrained("flax-community/swe-gpt-wiki")
# model = AutoModelWithLMHead.from_pretrained("flax-community/swe-gpt-wiki")

# generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
# set_seed(42)


#===========================================#
#              Streamlit Code               #
#===========================================#

# result = st.sidebar.selectbox(
#     "What language model would you like to try?",
#     ("Swedish", "Norwegian", "Danish", "Nordic"))

gpt_models = ["birgermoell/swedish-gpt", "flax-community/swe-gpt-wiki", "flax-community/nordic-gpt-wiki", "flax-community/norsk-gpt-wiki", "flax-community/dansk-gpt-wiki"]

selected_model = st.selectbox(
    "What language model would you like to try?",
    (gpt_models))

tokenizer = AutoTokenizer.from_pretrained(selected_model)
model = AutoModelWithLMHead.from_pretrained(selected_model)

generator = pipeline('text-generation', model=model, tokenizer=tokenizer)


st.title('GPT text generering')
desc = "Pröva GPT-modeller på flera språk. Skriv text nedanför för att utvärdera. Använder " + selected_model + " för att generera text. " + "Här kan du läsa mer om modellen https://huggingface.co/" + selected_model

st.write(desc)

num_sentences = st.number_input('Number of Characters', min_value=1, max_value=150, value=75)
user_input = st.text_input('Fyll i text att generera ifrån')


if st.button('Generate Text'):
    generated_text = generator(user_input, max_length=num_sentences, num_return_sequences=1)
    st.write(generated_text[0]["generated_text"])