Note_Retriever / app.py
ZacBl's picture
Upload 9 files
35bda59 verified
Raw
History Blame
1.73 kB
import streamlit as st
import numpy as np
from doc_preprocessing import process_files, get_embeddings
from vector_DB import VectorDatabase # Import the class
from llm_interaction import get_answer
# Initialize vector database (FAISS) - corrected instantiation
vector_database = VectorDatabase() #Instantiate the VectorDatabase Class
chunks_metadata = []
def main():
st.title("Document Query App")
uploaded_files = st.file_uploader(
"Upload PDF or Word files", accept_multiple_files=True, type=["pdf", "docx"]
)
query = st.text_input("Enter your query:")
if uploaded_files:
global chunks_metadata
all_chunks, all_embeddings, chunks_metadata = process_files(uploaded_files)
vector_database.add_data(all_embeddings, all_chunks, chunks_metadata) # use the method
st.session_state.files_processed = True
if query:
results = process_query(query)
display_results(results)
def process_query(query):
if vector_database.is_empty(): #Use the method
return "Please upload files first."
# query_embedding = get_embeddings([query])[0]
# results = vector_database.query(query_embedding, k=3) # use the method
query_embedding = get_embeddings([query])[0] # Get the embedding for the query
results = vector_database.query(query_embedding, k=3) # Get the top 2 results
return results
def display_results(results):
for result in results:
st.subheader("Answer")
st.subheader("Source")
st.write(f"File: {result['file_name']}, Chunk: {result['chunk_index']}")
st.subheader("Citations depuis le document :")
st.write(result["chunk_text"])
if __name__ == "__main__":
main()