Update app.py

app.py

@@ -1,124 +1,80 @@
-# Re-build-1
-
 import gradio as gr
-import faiss
-import numpy as np
-import os
-from datasets import load_dataset, Dataset
-from huggingface_hub import HfApi, hf_hub_download
-from PyPDF2 import PdfReader
 from transformers import AutoTokenizer, AutoModel
 import torch
+import numpy as np
+from PyPDF2 import PdfReader
+from pinecone import Pinecone, ServerlessSpec, CloudProvider, AwsRegion, VectorType
+import os
 
-# Set HF Dataset Name & Index File
-HF_DATASET_NAME = "aquibmoin/SCDD-Embeddings"
-INDEX_FILE = "faiss_index.faiss"
-
-# Load NASA Bi-Encoder
-bi_encoder_model_name = "nasa-impact/nasa-smd-ibm-st-v2"
-bi_tokenizer = AutoTokenizer.from_pretrained(bi_encoder_model_name)
-bi_model = AutoModel.from_pretrained(bi_encoder_model_name)
-
-# Initialize HF API
-hf_api = HfApi()
-
-# Function to extract text from a PDF
-def extract_text_from_pdf(pdf_file):
-    text = ""
-    with pdf_file as f:
-        reader = PdfReader(f)
-        for page in reader.pages:
-            text += page.extract_text() + "\n"
-    return text
-
-# Function to split text into chunks
-def get_chunks(text, chunk_size=500):
-    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
-
-# Function to generate embeddings
-def generate_embedding(text):
-    inputs = bi_tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    with torch.no_grad():
-        outputs = bi_model(**inputs)
-    embedding = outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
-    return embedding / np.linalg.norm(embedding)  # Normalize for FAISS
-
-# Function to load existing FAISS index from HF
-def load_existing_faiss_index():
-    try:
-        index_path = hf_hub_download(repo_id=HF_DATASET_NAME, filename=INDEX_FILE, repo_type="dataset")
-        index = faiss.read_index(index_path)
-        print("✅ Loaded existing FAISS index.")
-        return index
-    except:
-        print("⚠️ No existing FAISS index found. Creating a new one.")
-        return faiss.IndexFlatIP(768)
-
-# Main function to process PDFs & update HF Dataset
-def process_pdfs_and_store(pdf_files):
-    index = load_existing_faiss_index()
-
-    try:
-        dataset = load_dataset(HF_DATASET_NAME, split="train")
-        existing_chunks = dataset["chunk_text"]
-        existing_embeddings = [np.array(emb) for emb in dataset["embedding"]]
-        existing_files = dataset["source_file"]
-    except:
-        existing_chunks, existing_embeddings, existing_files = [], [], []
-    
-    all_chunks, all_embeddings = [], []
+# Load NASA-specific bi-encoder model
+tokenizer = AutoTokenizer.from_pretrained("nasa-impact/nasa-smd-ibm-st-v2")
+model = AutoModel.from_pretrained("nasa-impact/nasa-smd-ibm-st-v2")
+
+# Initialize Pinecone client
+pinecone_api_key = os.getenv('PINECONE_API_KEY')
+pc = Pinecone(api_key=pinecone_api_key)
+
+# Create Pinecone index if it doesn't exist
+index_name = "scdd-index"
+if index_name not in pc.list_indexes().names():
+    pc.create_index(
+        name=index_name,
+        dimension=768,
+        spec=ServerlessSpec(
+            cloud=CloudProvider.AWS,
+            region=AwsRegion.US_EAST_1
+        ),
+        vector_type=VectorType.DENSE,
+        metric="cosine"
+    )
 
+# Connect to the Pinecone index
+index = pc.Index(index_name)
+
+# Function to encode text using bi-encoder in batches
+def encode_chunks_batch(chunks, batch_size=8):
+    embeddings = []
+    for i in range(0, len(chunks), batch_size):
+        batch_chunks = chunks[i:i+batch_size]
+        inputs = tokenizer(batch_chunks, return_tensors='pt', padding=True, truncation=True, max_length=128)
+        with torch.no_grad():
+            output = model(**inputs)
+            batch_embeddings = output.last_hidden_state.mean(dim=1)
+            batch_embeddings = batch_embeddings / batch_embeddings.norm(dim=1, keepdim=True)
+            embeddings.extend(batch_embeddings.cpu().numpy())
+    return embeddings
+
+# Function to process PDFs and upsert embeddings to Pinecone
+def process_pdfs(pdf_files):
     for pdf_file in pdf_files:
-        text = extract_text_from_pdf(pdf_file)
-        chunks = get_chunks(text)
-        embeddings = [generate_embedding(chunk) for chunk in chunks]
-
-        all_chunks.extend(chunks)
-        all_embeddings.extend(embeddings)
-
-    all_embeddings_np = np.array(all_embeddings)
+        reader = PdfReader(pdf_file.name)
+        pdf_text = "".join(page.extract_text() for page in reader.pages if page.extract_text())
+        
+        # Split text into smaller chunks
+        chunks = [pdf_text[i:i+500] for i in range(0, len(pdf_text), 500)]
 
-    # Append new embeddings & chunks to the existing ones
-    combined_chunks = existing_chunks + all_chunks
-    combined_embeddings = existing_embeddings + list(all_embeddings_np)
-    combined_files = existing_files + [pdf_file.name for pdf_file in pdf_files for _ in range(len(all_chunks))]
+        # Generate embeddings in batches
+        embeddings = encode_chunks_batch(chunks, batch_size=8)
 
-    combined_embeddings_np = np.array(combined_embeddings)
+        # Prepare data for Pinecone
+        vectors = [
+            (f"{os.path.basename(pdf_file.name)}-chunk-{idx}", embedding.tolist(), {"text": chunk})
+            for idx, (embedding, chunk) in enumerate(zip(embeddings, chunks))
+        ]
 
-    # Update FAISS Index
-    index.add(all_embeddings_np)
+        # Upsert embeddings into Pinecone
+        index.upsert(vectors)
 
-    # Save & Upload Updated FAISS Index
-    faiss.write_index(index, INDEX_FILE)
-    hf_api.upload_file(path_or_fileobj=INDEX_FILE, path_in_repo=INDEX_FILE, repo_id=HF_DATASET_NAME, repo_type="dataset")
+    return f"Processed {len(pdf_files)} PDF(s) successfully and embeddings stored in Pinecone."
 
-    # Update & Push Dataset
-    dataset_dict = {
-        "chunk_text": combined_chunks,
-        "embedding": [emb.tolist() for emb in combined_embeddings_np],
-        "source_file": combined_files
-    }
+# Gradio Interface
+demo = gr.Interface(
+    fn=process_pdfs,
+    inputs=gr.Files(label="Upload PDFs", file_types=[".pdf"]),
+    outputs="text",
+    title="NASA Bi-encoder PDF Embedding & Pinecone Storage",
+    description="Upload PDF files to generate embeddings with NASA Bi-encoder and store in Pinecone."
+)
 
-    dataset = Dataset.from_dict(dataset_dict)
-    dataset.push_to_hub(HF_DATASET_NAME, split="train")
-
-    return f"✅ Successfully updated FAISS index & embeddings in {HF_DATASET_NAME}. Total Chunks: {len(combined_chunks)}."
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 SCDD Embeddings Generator - Hugging Face Spaces")
-    gr.Markdown("Upload PDFs to generate and store embeddings in a FAISS vector store on Hugging Face.")
-
-    pdf_input = gr.Files(file_types=[".pdf"], label="Upload Reference PDFs (Up to 3)", interactive=True)
-    submit_button = gr.Button("Generate & Store Embeddings")
-
-    output_text = gr.Textbox(label="Status")
-
-    submit_button.click(
-        fn=process_pdfs_and_store,
-        inputs=[pdf_input],
-        outputs=[output_text]
-    )
+demo.launch()
 
-# Launch Gradio App
-demo.launch(share=True)