pages/llama_custom_demo.py

import os

import openai
import streamlit as st

from models.llamaCustom import LlamaCustom
from utils.chatbox import chatbox

st.set_page_config(page_title="Llama", page_icon="🦙")

st.subheader("Llama Index with Custom LLM Demo")

if "messages" not in st.session_state:
    st.session_state.messages = []

if "openai_api_key" not in st.session_state:
    st.info("Enter your openai key to access the chatbot.")
else:
    option = st.selectbox(
        label="Select your model:", options=("bigscience/bloom-560m",)
    )

    with st.spinner("Initializing vector index"):
        model = LlamaCustom(model_name=option)

    chatbox("llama_custom", model)