app.py

import os

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

import streamlit as st
import weave
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_PROJECT = "paper_reader"

weave.init(f"{WANDB_PROJECT}")

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

with st.spinner('Loading the RAG pipeline...'):
    @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)