File size: 3,128 Bytes
4074cd7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | import streamlit as st
import fitz # PyMuPDF
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI
from dotenv import load_dotenv
import os
# Load environment variables
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
# Initialize OpenAI client
client = OpenAI(api_key=api_key)
# Utility Functions
def load_pdf(file):
"""Extract text from a PDF file."""
try:
doc = fitz.open(stream=file.read(), filetype="pdf")
return "".join([page.get_text() for page in doc])
except Exception as e:
st.error(f"Error reading PDF: {e}")
return ""
def split_text(text, chunk_size=1000, chunk_overlap=20):
"""Split text into manageable chunks."""
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, is_separator_regex=False
)
return text_splitter.create_documents([text])
def create_and_load_db(chunks, persist_directory="pdf_embeddings"):
"""Create and load ChromaDB."""
embeddings = HuggingFaceEmbeddings()
vectordb = Chroma.from_documents(documents=chunks, embedding=embeddings, persist_directory=persist_directory)
vectordb.persist()
return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
def generate_response(context, question):
"""Generate a response using OpenAI."""
try:
messages = [
{"role": "system", "content": "You are an assistant that answers questions based on PDF content."},
{"role": "user", "content": f"Context: {context}\n\nQuestion: {question}\n\nAnswer:"}
]
response = client.chat.completions.create(
model="gpt-3.5-turbo", # Replace with preferred model
messages=messages,
max_tokens=150,
)
return response.choices[0].message.content.strip()
except Exception as e:
st.error(f"Error generating response: {e}")
return None
# Main Application Logic
def main():
st.title("PDF Chatbot with OpenAI")
# Sidebar: File upload
uploaded_file = st.sidebar.file_uploader("Upload a PDF", type=['pdf'])
prompt = st.text_input("Ask a Question", "")
submitted = st.button("Submit")
if submitted and uploaded_file:
pdf_text = load_pdf(uploaded_file)
if pdf_text:
st.write("PDF Content Loaded!")
chunks = split_text(pdf_text)
vectordb = create_and_load_db(chunks)
if prompt:
docs = vectordb.similarity_search(prompt)
if docs:
context = docs[0].page_content
response = generate_response(context, prompt)
st.subheader("Generated Answer:")
st.write(response)
else:
st.warning("No relevant information found.")
else:
st.error("Unable to extract text from the PDF.")
if __name__ == "__main__":
main()
|