import streamlit as st
import gpt_neox

st.title('GPT-NeoX Streamlit App')

model = gpt_neox.GPTNeoX(model_name='D.A.N Rev III O.R.CA Edition')

@st.cache
def get_prediction(input_text):
    return model.predict(input_text)

st.write('Enter your input text below to get a prediction from the GPT-NeoX model:')
user_input = st.text_input('Input Text')

if user_input:
    prediction = get_prediction(user_input)
    st.write('Prediction:', prediction)