app.py

import os
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

import weave

import streamlit as st
from rag.rag import SimpleRAGPipeline

st.set_page_config(page_title="Chat with the Llama 3 paper!", page_icon="🦙", layout="centered", initial_sidebar_state="auto", menu_items=None)

wandb_api_key = st.sidebar.text_input('WANDB_API_KEY', type='password')
if len(wandb_api_key)>=10:
    os.environ["WANDB_API_KEY"] = wandb_api_key
else:
    st.stop()

WANDB_PROJECT = "paper_reader"

weave.init(f"{WANDB_PROJECT}")

st.title("Chat with the Llama 3 paper 💬🦙")

@st.cache_resource(show_spinner=False)
def load_rag_pipeline():
    rag_pipeline = SimpleRAGPipeline()
    rag_pipeline.build_query_engine()

    return rag_pipeline

if "rag_pipeline" not in st.session_state.keys():
    st.session_state.rag_pipeline = load_rag_pipeline()

rag_pipeline = st.session_state["rag_pipeline"]


def generate_response(query):
    response = rag_pipeline.predict(query)
    st.write_stream(response.response_gen)

with st.form('my_form'):
    query = st.text_area('Ask your question about the Llama 3 paper here:')
    submitted = st.form_submit_button('Submit')
    if submitted:
        generate_response(query)