File size: 3,008 Bytes
4aff460 f0422e2 4aff460 14e4af7 4aff460 14e4af7 4aff460 2ff3a9a |
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 |
import streamlit as st
import fitz # PyMuPDF
import faiss
from sentence_transformers import SentenceTransformer
import numpy as np
from phi.agent import Agent
from phi.model.groq import Groq
# Load embedding model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
def agent_response(question, retrieved_text):
agent = Agent(
model=Groq(id="llama-3.3-70b-versatile"),
markdown=True,
description="You are an AI assistant that provides the answer based on the provided document.",
instructions=[
f"First read the question carefully. The question is: **{question}**",
f"Then read the document provided to you as a text. The document is: \n**{retrieved_text}**\n",
"Finally answer the question based on the provided document only. Don't try to give random responses."
]
)
response = agent.run(question + '\n' + retrieved_text).content
return response
class PDFChatbot:
def __init__(self):
self.text_chunks = []
self.index = None
def process_pdf(self, pdf_file):
"""Extract text from PDF and create FAISS index."""
self.text_chunks = []
# Read the uploaded file into memory
with fitz.open("pdf", pdf_file.read()) as doc: # Use "pdf" as format
for page in doc:
self.text_chunks.append(page.get_text("text"))
# Embed the chunks
embeddings = embedding_model.encode(self.text_chunks, convert_to_numpy=True)
# Create FAISS index
self.index = faiss.IndexFlatL2(embeddings.shape[1])
self.index.add(embeddings)
return "PDF processed successfully!"
def chat(self, query):
"""Retrieve the most relevant chunk for a query."""
if self.index is None:
return "Please upload a PDF first."
# query_embedding = embedding_model.encode([query], convert_to_numpy=True)
# _, indices = self.index.search(query_embedding, 1) # Get top match
# retrieved_text = self.text_chunks[indices[0][0]]
# response = agent_response(query, retrieved_text)
query_embedding = embedding_model.encode([query], convert_to_numpy=True)
_, indices = self.index.search(query_embedding, 5) # Get top 5 matches
retrieved_texts = [self.text_chunks[idx] for idx in indices[0]]
retrieved_text_combined = "\n\n".join(retrieved_texts)
print('Retrieved_texts:',retrieved_text_combined)
response = agent_response(query, retrieved_text_combined)
return response
# Instantiate chatbot
chatbot = PDFChatbot()
st.title("Chat with your PDF")
uploaded_file = st.file_uploader("Upload your PDF", type=["pdf"])
if uploaded_file:
message = chatbot.process_pdf(uploaded_file)
st.success(message)
query = st.text_input("Ask a question")
if st.button("Ask"):
if query:
response = chatbot.chat(query)
st.markdown(f"**Answer:**\n\n{response}")
|