import streamlit as st
import requests
import time
from ast import literal_eval

    
def infer(
    prompt, 
    model_name, 
    max_new_tokens=10, 
    temperature=0.0, 
    top_p=1.0,
    top_k=40,
    num_completions=1,
    seed=42,
    stop="\n"
):

    model_name_map = {
        "GPT-JT-6B-v1": "Together-gpt-JT-6B-v1",
    }
    
    max_new_tokens = int(max_new_tokens)
    num_completions = int(num_completions)
    temperature = float(temperature)
    top_p = float(top_p)
    stop = stop.split(";")
    seed = seed
    
    assert 0 <= max_new_tokens <= 256
    assert 1 <= num_completions <= 5
    assert 0.0 <= temperature <= 10.0
    assert 0.0 <= top_p <= 1.0
    
    if temperature == 0.0:
        temperature = 1.0
        top_k = 1
    
    result = await st.session_state.together_web3.language_model_inference(
        from_dict(
            data_class=LanguageModelInferenceRequest,
            data={
                "model": model_name_map[model_name],
                "max_tokens": max_new_tokens,
                "prompt": prompt,
                "n": num_completions,
                "temperature": temperature,
                "top_k": top_k,
                "top_p": top_p,
                "stop": stop,
                "seed": seed,
                "echo": False,
            }
        ),
    )
    
    generated_text = result.choices[0].text
    
    for stop_word in stop:
        if stop_word in result:
            generated_text = generated_text[:generated_text.find(stop_word)]
    
    return generated_text

def set_preset():
    if st.session_state.preset == "Classification":
        
        st.session_state.prompt = '''Please classify the given sentence.
Possible labels:
1. <label_0>
2. <label_1>

Input: <sentence_0>
Label: <label_0>

Input: <sentence_1>
Label:'''
        st.session_state.temperature = "0.0"
        st.session_state.top_p = "1.0"
        
    elif st.session_state.preset == "Generation":
        
        st.session_state.prompt = '''Please write a story given keywords.

Input: bear, honey
Story:'''
        st.session_state.temperature = "1.0"
        st.session_state.top_p = "0.5"
    
    else:
        pass
    
    
def main():

    if 'preset' not in st.session_state:
        st.session_state.preset = "Classification"

    if 'prompt' not in st.session_state:
        st.session_state.prompt = "Please answer the following question:\n\nQuestion: Where is Zurich?\nAnswer:"

    if 'temperature' not in st.session_state:
        st.session_state.temperature = "0.0"

    if 'top_p' not in st.session_state:
        st.session_state.top_p = "1.0"

    if 'top_k' not in st.session_state:
        st.session_state.top_k = "40"
        
    if 'together_web3' not in st.session_state:
        st.session_state.together_web3 = TogetherWeb3()
        

    st.title("GPT-JT")

    col1, col2 = st.columns([1, 3])

    with col1:
        model_name = st.selectbox("Model", ["GPT-JT-6B-v1"])
        max_new_tokens = st.text_input('Max new tokens', "10")
        temperature = st.text_input('temperature', st.session_state.temperature)
        top_k = st.text_input('top_k', st.session_state.top_k)
        top_p = st.text_input('top_p', st.session_state.top_p)
        # num_completions = st.text_input('num_completions (only the best one will be returend)', "1")
        num_completions = "1"
        stop = st.text_input('stop, split by;', r'\n')
        # seed = st.text_input('seed', "42")
        seed = "42"

    with col2:

        preset = st.radio(
            "Recommended Configurations", 
            ('Classification', 'Generation'), 
            on_change=set_preset,
            key="preset",
            horizontal=True
        )

        prompt = st.text_area(
            "Prompt",
            value=st.session_state.prompt,
            max_chars=4096,
            height=400,
        )

        generated_area = st.empty()
        generated_area.text("(Generate here)")

        button_submit = st.button("Submit")

        if button_submit:
            generated_area.text(prompt)
            report_text = infer(
                prompt, model_name=model_name, max_new_tokens=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k,
                num_completions=num_completions, seed=seed, stop=literal_eval("'''"+stop+"'''"),
            )
            generated_area.text(prompt + report_text)

        
        
if __name__ == '__main__':
    main()