Spaces:
Sleeping
Sleeping
| import tempfile | |
| import os | |
| import streamlit as st | |
| from PyPDF2 import PdfReader | |
| from sentence_transformers import SentenceTransformer | |
| import faiss | |
| import numpy as np | |
| import time | |
| # Global variables for caching the model and embeddings | |
| model = None | |
| index = None | |
| embeddings = None | |
| text_chunks = [] | |
| # Function to process the uploaded PDF and save it temporarily | |
| def process_pdf(file): | |
| st.write("Processing uploaded PDF...") | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmpfile: | |
| tmpfile.write(file.read()) # Write the uploaded file's content to the temp file | |
| tmpfile_path = tmpfile.name # Get the temporary file path | |
| return tmpfile_path | |
| # Function to extract text from the PDF | |
| def extract_text_from_pdf(pdf_path): | |
| try: | |
| st.write("Extracting text from the PDF...") | |
| reader = PdfReader(pdf_path) | |
| text = "" | |
| for page in reader.pages: | |
| text += page.extract_text() | |
| return text | |
| except Exception as e: | |
| st.error(f"Error extracting text from PDF: {e}") | |
| return "" | |
| # Function to chunk text into smaller sections | |
| def chunk_text(text, chunk_size=200): | |
| try: | |
| st.write("Chunking text into smaller sections...") | |
| words = text.split() | |
| chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)] | |
| return chunks | |
| except Exception as e: | |
| st.error(f"Error chunking text: {e}") | |
| return [] | |
| # Function to load the embedding model | |
| def load_model(): | |
| global model | |
| st.write("Loading embedding model...") | |
| model = SentenceTransformer('all-MiniLM-L6-v2') | |
| # Function to generate embeddings | |
| def generate_embeddings(): | |
| global embeddings, text_chunks, index | |
| st.write("Generating embeddings...") | |
| embeddings = [] | |
| for chunk in text_chunks: | |
| embeddings.append(model.encode(chunk, convert_to_numpy=True)) | |
| embeddings = np.array(embeddings) | |
| # Build FAISS index | |
| st.write("Building FAISS index...") | |
| dimension = embeddings.shape[-1] | |
| index = faiss.IndexFlatL2(dimension) | |
| index.add(embeddings) | |
| # Main function to run the Streamlit app | |
| def main(): | |
| global embeddings, text_chunks, index, model | |
| st.title("PDF Embedding and Query System") | |
| # File uploader for the user to upload a PDF | |
| uploaded_file = st.file_uploader("Upload a PDF file", type=["pdf"]) | |
| if uploaded_file is not None: | |
| # Process the uploaded PDF and get its file path | |
| tmp_file_path = process_pdf(uploaded_file) | |
| # Extract text from the uploaded PDF | |
| pdf_text = extract_text_from_pdf(tmp_file_path) | |
| if not pdf_text: | |
| st.error("No text extracted from the PDF. Please upload a valid file.") | |
| return | |
| # Initialize Sentence-Transformer model and embeddings only once | |
| if model is None: | |
| load_model() | |
| # Chunk text into smaller sections for embedding generation | |
| if not text_chunks: | |
| text_chunks = chunk_text(pdf_text, chunk_size=200) | |
| # Generate embeddings only once | |
| if embeddings is None: | |
| generate_embeddings() | |
| # Query input field for users to enter their search queries | |
| query = st.text_input("Enter a query to search:") | |
| if query: | |
| # Generate embedding for the query | |
| query_embedding = model.encode([query], convert_to_numpy=True) | |
| # Perform similarity search using FAISS | |
| st.write("Searching...") | |
| start_time = time.time() | |
| D, I = index.search(query_embedding, k=5) | |
| end_time = time.time() | |
| # Display the results | |
| st.write(f"Query processed in {end_time - start_time:.2f} seconds.") | |
| for i in range(len(I[0])): | |
| st.write(f"Match {i + 1}: {text_chunks[I[0][i]]} (Distance: {D[0][i]:.4f})") | |
| if __name__ == "__main__": | |
| main() | |