LVKinyanjui/QueryYourDocs

1import streamlit as st2import pymupdf3import chromadb4from uuid import uuid45 6@st.cache_resource7def initdb():8    chroma_client = chromadb.Client()9    collection = chroma_client.get_or_create_collection(name="rag_collection")10    return collection11 12st.write("## Local RAG \n Get Insights from your documents")13 14file = st.file_uploader("Upload your Document Here to Query", type=['pdf'])15 16if file is not None:17    # Read file as bytes and save it.18    # PyMuPDF open can only load from file path19    bytes_data = file.getvalue()20    with open("data/uploaded_file.pdf", "wb") as fp:21        fp.write(bytes_data)22        doc = pymupdf.open(fp)23 24    texts = [str(page.get_text().encode("utf-8")) for page in doc]25 26    # VECTOR STORE27    collection = initdb()28 29    text_ids = [str(uuid4()) for text in texts]30    collection.add(documents=texts, ids=text_ids)31    st.write("Succesfully uploaded document to database.")32 33    # QUERY AREA34    query = st.text_input("Enter your query")35 36    if query:37        query_results = collection.query(38            query_texts=[query, ],39            n_results=5,40            include=["documents", ]41        )42 43        st.write("Database Query Matches")44        query_results45 46        # query_text = [" ".join([str(element) for element in inner_list])47        #     for inner_list in query_results["documents"]][0]48 49        # st.write("Database Query Matches")50        # st.markdown(query_text)51