googlecolabbuildvector(bgem3)

!pip install -q faiss-gpu-cu12 langchain-google-genai langchain-community pypdf PyMuPDF python-dotenv -q -U sentence-transformers

import os, time, json, zipfile, shutil, gc
import numpy as np
import faiss
import torch
from tqdm.auto import tqdm
from google.colab import drive, userdata
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI

# --- 1. SETUP & PATHS ---
print("🔌 Mounting Google Drive...")
drive.mount('/content/drive')

# CONFIGURATION
ZIP_SOURCE = "/content/singapore_data.zip"
INDEX_PATH = "/content/drive/MyDrive/sg_33k_laws.index"
META_PATH = "/content/drive/MyDrive/sg_33k_metadata.json"
TEMP_EXTRACT_DIR = "/content/temp_extracted_pdfs"
BATCH_SIZE = 20  # Save progress every 20 PDFs

# API KEY SETUP
try:
    API_KEY = userdata.get('GENAI_API_KEY')
except:
    API_KEY = input("Enter your Gemini API Key: ")
os.environ["GOOGLE_API_KEY"] = API_KEY

# --- 2. INITIALIZE MODELS ---
print("🚀 Loading BGE-M3 (Local GPU)...")

# FIX: Removed 'torch_dtype' to solve TypeError.
# The T4 GPU has enough VRAM (16GB) to run this model comfortably without it.
embedder = HuggingFaceEmbeddings(
    model_name="BAAI/bge-m3",
    model_kwargs={'device': 'cuda'},
    encode_kwargs={'normalize_embeddings': True, 'batch_size': 32}
)

llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)

# --- 3. ROBUST INDEX BUILDER ---
def build_or_resume_index():
    # A. EXTRACTION
    if not os.path.exists(TEMP_EXTRACT_DIR):
        print(f"📦 Extracting {ZIP_SOURCE} to local temp storage...")
        if not os.path.exists(ZIP_SOURCE):
            raise FileNotFoundError(f"❌ '{ZIP_SOURCE}' not found! Please upload 'singapore_data.zip' to your Google Drive.")

        with zipfile.ZipFile(ZIP_SOURCE, 'r') as z:
            z.extractall(TEMP_EXTRACT_DIR)

    # B. DISCOVERY
    all_pdfs = []
    for root, _, files in os.walk(TEMP_EXTRACT_DIR):
        for f in files:
            if f.lower().endswith(".pdf"):
                all_pdfs.append(os.path.join(root, f))

    print(f"🔍 Found {len(all_pdfs)} PDFs total.")

    # C. RESUME LOGIC
    metadata_map = []
    if os.path.exists(META_PATH):
        with open(META_PATH, 'r') as f:
            metadata_map = json.load(f)

        # Filter out PDFs we have already processed
        processed_sources = set([m['source'] for m in metadata_map])
        pdfs_to_process = [p for p in all_pdfs if os.path.basename(p) not in processed_sources]
        print(f"resuming skipping {len(processed_sources)} files. {len(pdfs_to_process)} remaining.")
    else:
        pdfs_to_process = all_pdfs
        print("Starting fresh index build...")

    if not pdfs_to_process:
        print("Index is up to date!")
        return

    # D. BATCH PROCESSING LOOP
    for i in range(0, len(pdfs_to_process), BATCH_SIZE):
        batch = pdfs_to_process[i:i+BATCH_SIZE]
        batch_texts = []
        batch_meta = []

        print(f"processing batch {i//BATCH_SIZE + 1}/{(len(pdfs_to_process)//BATCH_SIZE)+1}...")

        for path in tqdm(batch, leave=False):
            try:
                loader = PyPDFLoader(path)
                pages = loader.load()
                filename = os.path.basename(path)

                for page in pages:
                    clean_text = page.page_content[:2000]
                    if len(clean_text) < 50: continue

                    batch_texts.append(clean_text)
                    batch_meta.append({
                        "source": filename,
                        "page": page.metadata.get("page", 0) + 1,
                        "text": clean_text
                    })
            except Exception as e:
                print(f"fail to read {os.path.basename(path)}: {e}")

        if not batch_texts: continue

        # E. EMBEDDING & INDEXING
        embeddings = embedder.embed_documents(batch_texts)
        embeddings_np = np.array(embeddings).astype('float32')

        # Load existing index or create new one
        if os.path.exists(INDEX_PATH):
            index = faiss.read_index(INDEX_PATH)
        else:
            dim = embeddings_np.shape[1]
            index = faiss.IndexFlatL2(dim)

        index.add(embeddings_np)

        # F. SAVE CHECKPOINT
        faiss.write_index(index, INDEX_PATH)
        metadata_map.extend(batch_meta)
        with open(META_PATH, 'w') as f:
            json.dump(metadata_map, f)

        # G. MEMORY CLEANUP
        del embeddings, embeddings_np, batch_texts, batch_meta, index
        gc.collect()
        torch.cuda.empty_cache()

    print("indexing complete!")

# --- 4. CHAT SYSTEM ---
def start_chat():
    print("\n loading index from drive.")
    if not os.path.exists(INDEX_PATH):
        print("index not foun please run the builder first")
        return

    # Move index to GPU for fast search
    cpu_index = faiss.read_index(INDEX_PATH)
    res = faiss.StandardGpuResources()
    gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)

    with open(META_PATH, 'r') as f:
        metadata = json.load(f)

    print(f"✅ System Ready! ({len(metadata)} pages indexed)")
    print("--------------------------------------------------")

    while True:
        query = input("\nLawyerGPT (Type 'exit'): ")
        if query.lower() in ['exit', 'quit']: break

        # Search
        q_vec = np.array([embedder.embed_query(query)]).astype('float32')
        distances, indices = gpu_index.search(q_vec, k=5)

        context_parts = []
        seen_sources = set()

        for idx in indices[0]:
            if idx == -1: continue
            item = metadata[idx]
            key = f"{item['source']}_p{item['page']}"
            if key in seen_sources: continue
            seen_sources.add(key)

            context_parts.append(f"SOURCE: {item['source']} (Page {item['page']})\nCONTENT: {item['text']}")

        context_str = "\n\n".join(context_parts)

        system_prompt = f"""
        ROLE: Senior Singapore Legal Architect.
        TASK: Answer using ONLY the provided context.
        RULES:
        1. Cite every claim with [Source Name, Page X].
        2. If the answer is not in the context, say "Data not found in current index."
        3. Be extremely concise (Max 150 words).

        CONTEXT:
        {context_str}

        USER QUESTION: {query}
        """

        try:
            print(" analyzing.")
            response = llm.invoke(system_prompt)
            print(f"\n RESPONSE:\n{response.content}")
        except Exception as e:
            print(f" Error generating response: {e}")


build_or_resume_index()
#this just to test that vectors is not corrupted and #is working
# start chatting
start_chat()