ascii_art = r"""
.
                                             _..--------.._
   ________________________              _.-e        ||    e-._
 /'                    \ ||``\         /'    \\  |      |  |   `\
|:    ===               -||:  )      /'   \   _.--eeee--._     / `\
 \.____________________/_||../     /'  \    /e            e\ /  /  `\
    |       |                     ||    \ /'     \||  |/    `\       \
    |    || `\     ___-----------/e\  \  /    \           /   \  /    \
     \    \\  \.-ee   \  |  /   /|  | - |   \   .-eeee-.    /  \     - ;
     |           \  \     |  _/' |O |  |e\ __./'        `\   _  | - _  |
     |O    /____________..--e [] |__|- | [] ..   .---._   |  -  | _   _|
     |   ===____________ --|]=====__|  | [] ..  ( D>   )  |  _  |  _  =|
     |O    \  /    /    ee--. [] |  | -| []__    `---'e   |  _  |  -  e|
     |           /   /    |  `\  |O |  |_/   `\          /      | -    |
     /    //  .-.__   / |   \  `\|  | _ |   /  `-.____.-'   \  /    \  ;
    |    ||  /     eee------------\_/     \   /            \   /      /
   _|_______|______________       || /    \      /|  | |\    /  \    /
   _|_______|______________       || /    \      /|  | |\    /  \    /
 /'                    \ ||``\     \   / / `\_             /'  \    /
|:   ===                -||:  )     `\        e-__    __.-e \  \  /'
 \.____________________/_||../        `\  /   /   eeee |        /'
                                        `-._     |  ||    \ _.-'
                                            e-..________..-e

This ASCII pic can be found at
https://asciiart.website/index.php?art=television/star%20trek
"""

from model import LLM
import streamlit as st
import streamlit_scrollable_textbox as stx

st.text(ascii_art)

# with st.spinner("Please wait... loading model"):
llm = LLM()

demo_text = """DATA: The ship has gone into warp, sir.
RIKER: Who gave the command?
DATA: Apparently no one. Helm and navigation controls are not functioning. Our speed is now warp seven-point-three and holding.
PICARD: Picard to Engineering. Mister La Forge, what's going on down there?
"""
text = st.text_area("First few lines of the script:", demo_text, height=300)

col1, col2 = st.columns(2)

with col1:
    temp = st.slider('Temperature', 0.0, 1.0, 1.0, 0.1)
    max_len = st.number_input('Max new tokens', min_value=1, max_value=2048, value=128)

with col2:
    top_p = st.slider('p', 0.0, 1.0, 0.95, 0.1)
    top_k = st.slider('k', 1, 100, 2, 5)

if st.button("Generate"):
    with st.spinner("Generating text..."):
        stx.scrollableTextbox(llm.generate(text, max_len, temp, top_k, top_p), height=400)