Update data_loader.py

data_loader.py

@@ -1,49 +1,52 @@
-import os
-import zipfile
-import pdfplumber
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain.vectorstores import FAISS
-
-DATASET_ZIP = "dataset.zip"
-FAISS_INDEX_FILE = "faiss_index"
-
-def extract_dataset():
-    """Extract dataset.zip contents if not already extracted."""
-    if os.path.exists(DATASET_ZIP):
-        with zipfile.ZipFile(DATASET_ZIP, 'r') as zip_ref:
-            zip_ref.extractall("./")
-            print("✅ Dataset extracted!")
-    else:
-        print("⚠️ No dataset.zip found, ensure financial data is available.")
-
-def extract_text_from_pdfs():
-    """Extract text from all PDFs in root directory."""
-    text_data = ""
-    for file in os.listdir("./"):
-        if file.endswith(".pdf"):
-            with pdfplumber.open(file) as pdf:
-                for page in pdf.pages:
-                    text_data += page.extract_text() or ""
-    return text_data
-
-def create_vector_store(api_key):
-    """Create FAISS vector database from extracted text."""
-    extract_dataset()
-    text_data = extract_text_from_pdfs()
-    
-    if not text_data:
-        print("⚠️ No valid text extracted.")
-        return
-    
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-    text_chunks = text_splitter.split_text(text_data)
-
-    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
-    vector_db = FAISS.from_texts(text_chunks, embedding=embeddings)
-    vector_db.save_local(FAISS_INDEX_FILE)
-    print("✅ FAISS index created and saved!")
-
-if __name__ == "__main__":
-    api_key = input("Enter Google API Key: ")
-    create_vector_store(api_key)
+import os
+import pdfplumber
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+
+DATASET_DIR = "."  # Root directory (Hugging Face doesn't allow separate dataset folders)
+FAISS_INDEX_PATH = "financial_faiss_index"
+
+def get_pdf_text(pdf_files):
+    """Extracts text from PDFs."""
+    text = ""
+    for pdf in pdf_files:
+        with pdfplumber.open(pdf) as reader:
+            for page in reader.pages:
+                text += page.extract_text() or ""  # Handle NoneType
+    return text.strip()
+
+def preprocess_and_store_embeddings(api_key):
+    """Extracts text from financial documents, creates embeddings, and saves FAISS index."""
+    financial_text = ""
+    
+    # Process all PDFs in the root directory
+    for file in os.listdir(DATASET_DIR):
+        if file.endswith(".pdf"):
+            file_path = os.path.join(DATASET_DIR, file)
+            financial_text += get_pdf_text([file_path]) + "\n\n"
+
+    if not financial_text:
+        print("No financial documents found. Please upload PDFs.")
+        return False
+
+    # Split text into chunks
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=10000, chunk_overlap=1000)
+    text_chunks = text_splitter.split_text(financial_text)
+
+    # Generate embeddings
+    embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=api_key)
+    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
+
+    # Save FAISS index in root folder
+    vector_store.save_local(FAISS_INDEX_PATH)
+    print("✅ FAISS index saved successfully!")
+
+    return True
+
+if __name__ == "__main__":
+    api_key = os.getenv("GOOGLE_API_KEY")
+    if api_key:
+        preprocess_and_store_embeddings(api_key)
+    else:
+        print("❌ Google API Key not found. Please provide a valid key.")