khalil
Create app.py
317cf00 verified
Raw
History Blame Contribute Delete
2.63 kB
import os
import fitz # PyMuPDF for PDF processing
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import streamlit as st
from groq import Groq
from tempfile import NamedTemporaryFile
# Set up the Groq client with your API key
client = Groq(api_key="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")
# Step 1: Function to extract text from PDF
def extract_text_from_pdf(pdf_path):
doc = fitz.open(pdf_path)
text = ""
for page in doc:
text += page.get_text()
doc.close()
return text
# Step 2: Function to split extracted text into chunks for retrieval
def chunk_text(text, chunk_size=1000):
words = text.split()
chunks = []
for i in range(0, len(words), chunk_size):
chunk = " ".join(words[i:i+chunk_size])
chunks.append(chunk)
return chunks
# Step 3: Retrieve the most relevant chunk using TF-IDF and cosine similarity
def retrieve_chunk(question, chunks):
vectorizer = TfidfVectorizer().fit_transform([question] + chunks)
question_vector = vectorizer[0]
chunk_vectors = vectorizer[1:]
similarities = cosine_similarity(question_vector, chunk_vectors).flatten()
best_chunk_index = np.argmax(similarities)
return chunks[best_chunk_index]
# Step 4: Generate an answer using the Groq API's language model
def generate_answer(retrieved_text, question):
prompt = f"Based on the following text, answer the question:\n\nText: {retrieved_text}\n\nQuestion: {question}"
chat_completion = client.chat.completions.create(
messages=[{"role": "user", "content": prompt}],
model="llama3-8b-8192"
)
return chat_completion.choices[0].message.content
# Step 5: Streamlit UI for PDF upload and Q&A
def main():
st.title("PDF Question-Answer Chatbot")
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
if uploaded_file:
with NamedTemporaryFile(delete=False) as tmp_file:
tmp_file.write(uploaded_file.getvalue())
pdf_path = tmp_file.name
# Extract text from the uploaded PDF and chunk it
text = extract_text_from_pdf(pdf_path)
chunks = chunk_text(text)
question = st.text_input("Ask a question:")
if st.button("Get Answer"):
if question:
retrieved_text = retrieve_chunk(question, chunks)
answer = generate_answer(retrieved_text, question)
st.write("Answer:", answer)
else:
st.write("Please enter a question.")
if __name__ == "__main__":
main()