Rag-raptor.py

import streamlit as st
from st_chat_message import message
from dotenv import load_dotenv
import fitz

load_dotenv()

from raptor.raptor import RetrievalAugmentation # noqa

RA = RetrievalAugmentation()


def upload_document():
    uploaded_file = st.file_uploader("Upload Document", type=['pdf', 'txt'])

    if uploaded_file:
        st.success("Document uploaded successfully", icon="😎")

        file_contents = ""

        if uploaded_file.type == 'text/plain':
            file_contents = uploaded_file.read().decode('utf-8')
        elif uploaded_file.type == 'application/pdf':
            with fitz.open("pdf", uploaded_file.getvalue()) as doc:
                texts = [page.get_text() for page in doc]
            file_contents = " ".join(texts)

        RA.add_documents(file_contents)
        st.spinner("Creating your document index...")
        chat_interface()


def chat_interface():
    message("I am your AI Assistant. Ask any question about the document you uploaded !")
    user_input = st.chat_input("Enter your message here")

    if user_input:
        message(user_input, is_user=True)
        answer = RA.answer_question(question=user_input)
        message(answer)


def main():
    st.title("RAG Chat App Using Recursive Abstractive Processing for Tree-Organized Retrieval")
    upload_document()


if __name__ == "__main__":
    main()