from contextlib import nullcontext
import torch
import tiktoken
from model import GPTConfig, GPT
import streamlit as st

# -----------------------------------------------------------------------------
st.set_page_config(page_title="Translation Demo", page_icon=":milky_way:", layout="wide")
st.subheader("Gerador Canções de musica brasileira")

# ----------------------------------------------
max_new_tokens = 200 # number of tokens generated in each sample
temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
top_k = 200 
seed = 1337
device = 'cpu' 
dtype = 'bfloat16' 
# -----------------------------------------------------------------------------

ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
ctx = nullcontext() 
checkpoint = torch.load('ckpt.pt', map_location='cpu')
gptconf = GPTConfig(**checkpoint['model_args'])
model = GPT(gptconf)
state_dict = checkpoint['model']
model.load_state_dict(state_dict)



def gera_texto(start, temperature, max_new_tokens, seed, num_samples):
    torch.manual_seed(seed)
    enc = tiktoken.get_encoding("gpt2")
    encode = lambda s: enc.encode(s, allowed_special={"\n"})
    decode = lambda l: enc.decode(l)


    # encode the beginning of the prompt
    start_ids = encode(start)
    x = (torch.tensor(start_ids, dtype=torch.long, device='cpu')[None, ...])

    # run generation
    geracoes = ""
    with torch.no_grad():
        with ctx:
            for k in range(num_samples):
                y = model.generate(x, max_new_tokens, temperature=temperature, top_k=top_k)
                geracoes =  decode(y[0].tolist())
                st.text_area("Gerado {}".format(k+1), value= geracoes, height=300, placeholder="")
                
with st.form("my_form"):
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        int_samples = st.slider('Exemplos', min_value=1, max_value=10, value=5, step=1)
    with col2:
        int_seed = st.slider('Seed', min_value=1, max_value=1500, value=1337, step=1)
    with col3:
        int_size = st.slider('Num Tokens', min_value=20, max_value=500, value=160, step=5)
    with col4:
        int_temp = st.number_input("Temperatura",min_value=0.1,max_value=2.0,value=0.8,step=0.1,format="%.1f")

    source = st.text_area("Escolha uma frase inicial", value="Voce e tao linda", placeholder="Entre com o inicio da musica...")

    submitted = st.form_submit_button("Gerar músicas")
    if submitted:
        with st.spinner("Gerando exemplos ..."):
            gera_texto(source,int_temp,int_size,int_seed, int_samples)
        
st.write("8 milhões de tokens, 16 camadas de atenção. Três dias de treinamento from stratch")
st.write("A preparação dos dados demorou um longo final de semana.")
st.write("Agradecimentos ao [Gabriel](https://www.linkedin.com/in/go2035/) pela ajuda no scrap.")
st.markdown("""---""")
original_title = '<p style="font-family:Verdana; color:Blue; font-size: 12px;">Gosta de IA ou é um maker por natureza ? Conecte-se ao meu <a href=https://www.linkedin.com/in/israeloliveira2035/> linkedin</a> e vamos conversar  !</p>'
st.markdown(original_title, unsafe_allow_html=True)
st.write("Made with [nanoGPT](https://github.com/karpathy/nanoGPT) e [ColabPro+](https://colab.research.google.com/signup)")