|
import os |
|
import shutil |
|
import streamlit as st |
|
from llm import load_llm, response_generator |
|
from vector_store import load_vector_store, process_pdf |
|
from uuid import uuid4 |
|
|
|
|
|
|
|
repo_id = "Qwen/Qwen2.5-3B-Instruct-GGUF" |
|
filename = "qwen2.5-3b-instruct-q5_k_m.gguf" |
|
|
|
|
|
llm = load_llm(repo_id, filename) |
|
vector_store = load_vector_store() |
|
|
|
|
|
st.title("PDF QA") |
|
|
|
if "messages" not in st.session_state: |
|
vector_store.reset_collection() |
|
if os.path.exists("./temp"): |
|
shutil.rmtree("./temp") |
|
st.session_state.messages = [] |
|
|
|
|
|
for message in st.session_state.messages: |
|
with st.chat_message(message["role"]): |
|
if message["role"] == "user": |
|
st.write(message["content"]) |
|
else: |
|
st.write(message["content"]) |
|
|
|
|
|
if prompt := st.chat_input("What is up?"): |
|
|
|
st.session_state.messages.append({"role": "user", "content": prompt}) |
|
|
|
with st.chat_message("user"): |
|
st.markdown(prompt) |
|
|
|
|
|
with st.chat_message("assistant"): |
|
retriever = vector_store.as_retriever(search_kwargs={"k": 3}) |
|
response = response_generator(llm, st.session_state.messages, prompt, retriever) |
|
|
|
st.markdown(response["answer"]) |
|
with st.expander("See context"): |
|
st.write(response["context"]) |
|
|
|
|
|
st.session_state.messages.append( |
|
{"role": "assistant", "content": response["answer"]} |
|
) |
|
|
|
with st.sidebar: |
|
st.title("PDFs") |
|
st.write("Upload your pdfs here") |
|
uploaded_files = st.file_uploader( |
|
"Choose a PDF file", accept_multiple_files=True, type="pdf" |
|
) |
|
if uploaded_files is not None: |
|
st.session_state.uploaded_pdf = True |
|
for uploaded_file in uploaded_files: |
|
temp_dir = "./temp" |
|
if not os.path.exists(temp_dir): |
|
os.makedirs(temp_dir) |
|
temp_file = f"./temp/{uploaded_file.name}-{uuid4()}.pdf" |
|
with open(temp_file, "wb") as file: |
|
file.write(uploaded_file.getvalue()) |
|
|
|
st.write("filename:", uploaded_file.name) |
|
process_pdf(temp_file, vector_store) |
|
st.success("PDFs processed successfully. ✅") |
|
|