File size: 2,643 Bytes
5c72b8f |
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 |
import os
import streamlit as st
import fitz # PyMuPDF
import faiss
import numpy as np
import pickle
from sentence_transformers import SentenceTransformer
import tiktoken
from groq import Groq
# Initialize embedding model
embed_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
text = "\n".join([page.get_text("text") for page in doc])
return text
# Function to split text into chunks
def chunk_text(text, chunk_size=512):
tokenizer = tiktoken.get_encoding("cl100k_base")
tokens = tokenizer.encode(text)
chunks = [tokens[i:i+chunk_size] for i in range(0, len(tokens), chunk_size)]
return ["".join(tokenizer.decode(chunk)) for chunk in chunks]
# Function to generate embeddings
def generate_embeddings(chunks):
return embed_model.encode(chunks, convert_to_numpy=True)
# Function to store embeddings in FAISS
def store_in_faiss(embeddings, chunks):
dimension = embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(embeddings)
with open("faiss_index.pkl", "wb") as f:
pickle.dump((index, chunks), f)
return index
# Function to load FAISS index
def load_faiss():
with open("faiss_index.pkl", "rb") as f:
index, chunks = pickle.load(f)
return index, chunks
# Function to search FAISS
def search_faiss(query, top_k=3):
query_embedding = embed_model.encode([query])
index, chunks = load_faiss()
_, indices = index.search(query_embedding, top_k)
results = [chunks[i] for i in indices[0]]
return results
# Function to interact with Groq API
def query_groq(query):
client = Groq(api_key=os.getenv("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"))
response = client.chat.completions.create(
messages=[{"role": "user", "content": query}],
model="llama-3.3-70b-versatile"
)
return response.choices[0].message.content
# Streamlit UI
st.title("RAG-based PDF Q&A App")
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
if uploaded_file:
st.write("Processing PDF...")
text = extract_text_from_pdf(uploaded_file)
chunks = chunk_text(text)
embeddings = generate_embeddings(chunks)
store_in_faiss(embeddings, chunks)
st.success("PDF processed and indexed!")
query = st.text_input("Ask a question:")
if query:
retrieved_chunks = search_faiss(query)
context = " ".join(retrieved_chunks)
response = query_groq(f"Context: {context} \n Question: {query}")
st.write("### Answer:")
st.write(response)
|