app.py

import streamlit as st
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from transformers import AutoModelForCausalLM, AutoTokenizer
import tempfile
import os

st.set_page_config(page_title="Document QA Bot")

if "vector_store" not in st.session_state:
    st.session_state.vector_store = None

def process_text(text):
    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = splitter.create_documents([text])
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    return FAISS.from_documents(chunks, embeddings)

def process_pdf(file):
    with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
        tmp_file.write(file.getvalue())
        loader = PyPDFLoader(tmp_file.name)
        pages = loader.load()
        os.unlink(tmp_file.name)
    return process_text('\n'.join(page.page_content for page in pages))

st.title("Document QA Bot")

uploaded_file = st.file_uploader("Upload Document", type=["txt", "pdf"])
if uploaded_file:
    with st.spinner("Processing document..."):
        if uploaded_file.type == "text/plain":
            text = uploaded_file.getvalue().decode()
            st.session_state.vector_store = process_text(text)
        else:
            st.session_state.vector_store = process_pdf(uploaded_file)
        st.success("Document processed!")

if st.session_state.vector_store:
    if question := st.chat_input("Ask a question about the document:"):
        results = st.session_state.vector_store.similarity_search(question)
        context = "\n".join(doc.page_content for doc in results)
        st.chat_message("user").write(question)
        st.chat_message("assistant").write(context)