app.py

ascii_art = r"""
      _               _            _    
  ___| |_ __ _ _ __  | |_ _ __ ___| | __
 / __| __/ _` | '__| | __| '__/ _ \ |/ /
 \__ \ || (_| | |    | |_| | |  __/   < 
 |___/\__\__,_|_|     \__|_|  \___|_|\_\

"""

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)

col1, col2 = st.columns(2)

with col1:
    temp = st.slider('Temperature', 0, 1, 1)
    max_len = st.number_input('Max length', min_value=1, max_value=2048, value=512)

with col2:
    top_p = st.slider('p', 0, 1, 0.95)
    top_k = st.slider('k', 1, 100, 50)

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