import streamlit as st
from init import init

init()

st.header("Settings")

st.session_state.system_prompt = st.text_input("Input", st.session_state.system_prompt)

st.session_state.params["model"] = st.selectbox(
    "Select Model", ["llama3.2", "falcon-mamba", "mistral-nemo"]
)

st.session_state.params["temperature"] = st.slider(
    "temperature",
    min_value=0.0,
    max_value=1.0,
    value=0.8,
)

st.session_state.params["max_tokens"] = st.select_slider(
    label="max_tokens",
    options=[16, 64, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072],
    value=65536,
)
st.session_state.params["top_p"] = st.slider(
    label="top_p", min_value=0.01, max_value=1.00, value=0.95
)
st.session_state.params["min_p"] = st.slider(
    label="min_p", min_value=0.01, max_value=1.00, value=0.05
)
st.session_state.params["typical_p"] = st.slider(
    label="typical_p", min_value=0.0, max_value=1.0, value=1.0
)
st.session_state.params["frequency_penalty"] = st.slider(
    label="frequency_penalty", min_value=0.0, max_value=1.0, value=0.0
)
st.session_state.params["presence_penalty"] = st.slider(
    label="presence_penalty", min_value=0.0, max_value=1.0, value=0.0
)
st.session_state.params["repeat_penalty"] = st.slider(
    label="repeat_penalty", min_value=0.0, max_value=1.0, value=1.0
)
st.session_state.params["top_k"] = st.slider(
    label="top_k", min_value=1, max_value=100, value=40
)
st.session_state.params["tfs_z"] = st.slider(
    label="tfs_z", min_value=0.0, max_value=1.0, value=1.0
)

mirostat_mode = ["disabled", "Mirostat", "Mirostat 2.0"]


def label_mirostat_mode(i):
    return mirostat_mode[i]


st.session_state.params["mirostat_mode"] = st.selectbox(
    label="mirostat_mode",
    options=[0, 1, 2],
    format_func=label_mirostat_mode,
)
st.session_state.params["mirostat_tau"] = st.slider(
    label="mirostat_tau", min_value=0.0, max_value=10.0, value=5.0
)
st.session_state.params["mirostat_eta"] = st.slider(
    label="mirostat_eta", min_value=0.01, max_value=1.0, value=0.1
)