import streamlit as st
from main import *

# Set the title of the web app
st.title("Text and Token Input App")

# Input for the user's prompt
prompt = st.text_area("Enter a prompt above 15 words:")

# Input for the number of tokens to be generated
num_tokens = st.number_input("Enter the number of tokens to be generated:", min_value=1, step=1)


# Function to generate output based on the prompt and number of tokens
def generate(prompt, num_tokens):
    output = gen(prompt, num_tokens)
    return output


# Button to perform the calculation
if st.button("Generate Output"):
    if len(prompt.split()) < 9:
        st.warning("Please enter a prompt with more than 15 words.")
    else:
        output = generate(prompt, num_tokens)

        # Display prompt length and sum of prompt length and number of tokens
        prompt_length = len(prompt.split())
        total_sum = prompt_length + num_tokens

        # Display output in a box below or next to the input prompts


        st.subheader("Generated Output")
        st.write(output)