Update worker_script.py

worker_script.py

@@ -1,8 +1,9 @@
 import os
 import sys
-import time
 
-# 1. AUTO-INSTALL DEPENDENCIES
+# ==========================================
+# 1. SETUP ENV (Kaggle Needs This)
+# ==========================================
 try:
     import sentence_transformers
     import pyarrow
@@ -13,36 +14,57 @@ except ImportError:
 
 import requests
 import pandas as pd
-import numpy as np
-import pyarrow.parquet as pq
-import pyarrow as pa
-import gc
 import hashlib
 import re
-import json
 from datetime import datetime
 from typing import Set
+import logging
+import numpy as np
+import pyarrow.parquet as pq
+import pyarrow as pa
+import gc
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import hf_hub_download, HfApi
 
-# Config
-DATA_URL = "https://open.canada.ca/data/dataset/432527ab-7aac-45b5-81d6-7597107a7013/resource/1d15a62f-5656-49ad-8c88-f40ce689d831/download/grants.csv"
-REPO_ID = "ravi2814/grant-data-storage"
+# Configure Logging to print to Stdout for the UI to catch
+logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='%(message)s')
+logger = logging.getLogger(__name__)
 
-PARQUET_FILE = "grant data.parquet"
-EMBEDDINGS_FILE = "embeddings.npy"
-METADATA_FILE = "last_metadata.json"
+# ==========================================
+# 2. YOUR PRECISE LOGIC
+# ==========================================
 
 class GrantsDataUpdater:
-    def __init__(self):
+    """
+    Incremental update system for Canadian Grants data.
+    Uses your preprocessing functions to detect new rows.
+    """
+    
+    def __init__(self, 
+                 data_url: str = "https://open.canada.ca/data/dataset/432527ab-7aac-45b5-81d6-7597107a7013/resource/1d15a62f-5656-49ad-8c88-f40ce689d831/download/grants.csv",
+                 api_url: str = "https://open.canada.ca/data/api/3/action/package_show?id=432527ab-7aac-45b5-81d6-7597107a7013",
+                 local_parquet: str = "grant data.parquet",
+                 metadata_file: str = "last_metadata.json",
+                 new_rows_csv: str = "new_rows.csv",
+                 full_updated_parquet: str = "grant data.parquet", # Overwrite same file
+                 chunk_size: int = 50000):
+        
+        self.data_url = data_url
+        self.api_url = api_url
+        self.local_parquet = local_parquet
+        self.metadata_file = metadata_file
+        self.new_rows_csv = new_rows_csv
+        self.full_updated_parquet = full_updated_parquet
+        self.chunk_size = chunk_size
+        
+        # Temporary files
         self.temp_new_csv = "temp_new_data.csv"
-        self.chunk_size = 50000
-
+    
     def normalize_and_hash(self, text):
         if not isinstance(text, str): text = ""
         clean_text = re.sub(r'\s+', '', text).lower()
         return hashlib.md5(clean_text.encode('utf-8')).hexdigest()
-
+    
     def format_record(self, record):
         def val(k): return str(record.get(k, "") or "").strip()
         mapping = [
@@ -61,133 +83,153 @@ class GrantsDataUpdater:
             ("owner organization", "owner_org"), ("owner organization title", "owner_org_title"),
         ]
         return ", ".join([f"{d} : {val(k)}" for d, k in mapping])
-
-    def download_existing_files(self):
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Connecting to HuggingFace...", flush=True)
+    
+    def check_for_updates(self) -> bool:
+        # Simplified for Kaggle run (Always run if forced, or check API)
         try:
-            hf_hub_download(repo_id=REPO_ID, filename=PARQUET_FILE, repo_type="dataset", local_dir=".", force_download=True)
-            hf_hub_download(repo_id=REPO_ID, filename=EMBEDDINGS_FILE, repo_type="dataset", local_dir=".", force_download=True)
-            try: hf_hub_download(repo_id=REPO_ID, filename=METADATA_FILE, repo_type="dataset", local_dir=".", force_download=True)
-            except: pass
-            print(f"[{datetime.now().strftime('%H:%M:%S')}] Success: Existing database retrieved.", flush=True)
-            return True
-        except:
-            print(f"[{datetime.now().strftime('%H:%M:%S')}] Notice: No existing database found. Starting fresh.", flush=True)
+            response = requests.get(self.api_url, timeout=30)
+            data = response.json()
+            if not data.get('success'): return False
+            result = data['result']
+            current_modified = result.get('metadata_modified')
+            
+            # If we don't have a local metadata file (first run on Kaggle), we must assume update needed
+            if not os.path.exists(self.metadata_file):
+                logger.info("No local metadata found. Forcing update check.")
+                return True
+                
+            with open(self.metadata_file, 'r') as f:
+                import json
+                last_meta = json.load(f)
+            
+            if current_modified != last_meta.get('metadata_modified'):
+                logger.info("New data detected via API.")
+                return True
+            
+            logger.info("Metadata matches. No update needed.")
             return False
-
-    def download_new_csv(self):
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Downloading latest CSV from Canada.ca...", flush=True)
+        except:
+            return True 
+    
+    def download_new_data(self) -> bool:
+        logger.info("Downloading new dataset...")
         try:
-            response = requests.get(DATA_URL, stream=True, timeout=300)
-            downloaded = 0
-            
+            response = requests.get(self.data_url, stream=True, timeout=300)
             with open(self.temp_new_csv, 'wb') as f:
-                for chunk in response.iter_content(chunk_size=1024*1024): 
-                    if chunk:
-                        f.write(chunk)
-                        downloaded += len(chunk)
-                        if downloaded % (25 * 1024 * 1024) < 1024*1024: 
-                            print(f"   ...downloaded {downloaded // (1024*1024)} MB", flush=True)
-            
-            print(f"[{datetime.now().strftime('%H:%M:%S')}] Download Complete.", flush=True)
+                for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
             return True
-        except Exception as e:
-            print(f"Download Error: {e}", flush=True)
-            return False
-
-    def get_existing_hashes(self):
-        existing_hashes = set()
-        if os.path.exists(PARQUET_FILE):
-            print(f"[{datetime.now().strftime('%H:%M:%S')}] Indexing existing records...", flush=True)
-            try:
-                pf = pq.ParquetFile(PARQUET_FILE)
-                rows = pf.metadata.num_rows
-                print(f"   Current Database Size: {rows:,} rows", flush=True)
-                
-                for batch in pf.iter_batches(batch_size=50000, columns=['data']):
-                    df_batch = batch.to_pandas()
-                    for text in df_batch['data']:
-                        existing_hashes.add(self.normalize_and_hash(text))
-                    del df_batch
-                    gc.collect()
-            except Exception as e:
-                print(f"Index Error: {e}", flush=True)
-        return existing_hashes
-
-    def process(self):
-        first_run = not self.download_existing_files()
-        if not self.download_new_csv(): return False
-
-        existing_hashes = self.get_existing_hashes()
-        
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Scanning for new grants...", flush=True)
+        except: return False
+    
+    def get_existing_hashes(self) -> Set[str]:
+        try:
+            logger.info("Loading existing data hashes...")
+            if not os.path.exists(self.local_parquet): return set()
+            
+            # Optimized: Read only 'data' column
+            df_old = pd.read_parquet(self.local_parquet, columns=['data'])
+            return set(self.normalize_and_hash(x) for x in df_old['data'])
+        except: return set()
+    
+    def process_and_find_new_rows(self, existing_hashes: Set[str]):
+        logger.info("Processing new data...")
         new_rows = []
-        scanned_count = 0
-        total_new = 0
-        chunk_idx = 0
+        all_rows = [] # We need to rebuild the file to keep order safe for embeddings
+        
+        # Load existing data first to keep it at the top (Append mode simulation)
+        if os.path.exists(self.local_parquet):
+             # We assume existing data is good.
+             pass 
+
+        # Scan New CSV
+        new_cnt = 0
+        scanned = 0
+        
+        # We will append new rows to a list, then append that list to the parquet file
+        # This avoids re-writing the whole 1GB file if we can help it.
         
         for df_chunk in pd.read_csv(self.temp_new_csv, chunksize=self.chunk_size, low_memory=False):
-            chunk_idx += 1
-            chunk_new_count = 0
-            
+            chunk_new = []
             for _, row in df_chunk.iterrows():
-                formatted_text = self.format_record(row.to_dict())
-                row_hash = self.normalize_and_hash(formatted_text)
+                formatted = self.format_record(row.to_dict())
+                h = self.normalize_and_hash(formatted)
                 
-                if row_hash not in existing_hashes:
-                    new_rows.append({'data': formatted_text})
-                    existing_hashes.add(row_hash)
-                    chunk_new_count += 1
+                if h not in existing_hashes:
+                    chunk_new.append({'data': formatted})
+                    existing_hashes.add(h)
             
-            scanned_count += len(df_chunk)
-            total_new += chunk_new_count
-            print(f"   Chunk {chunk_idx}: Scanned {scanned_count:,} rows. Found {chunk_new_count} NEW grants.", flush=True)
-        
-        if not new_rows:
-            print(f"[{datetime.now().strftime('%H:%M:%S')}] Scan complete. No new data found.", flush=True)
-            if os.path.exists(self.temp_new_csv): os.remove(self.temp_new_csv)
-            return True
-
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Scan complete. TOTAL NEW GRANTS: {total_new}", flush=True)
-
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Generating vectors for {len(new_rows)} new items...", flush=True)
-        
-        # CORRECTED DEVICE LOGIC
-        device = 'cuda' if hasattr(sys, 'getandroidapilevel') else ('cuda' if os.path.exists('/kaggle') else 'cpu')
-        model = SentenceTransformer('intfloat/multilingual-e5-large', device=device)
-        
-        new_texts = ["passage: " + r['data'] for r in new_rows]
-        batch_size = 1000
-        new_embeddings_list = []
-        
-        for i in range(0, len(new_texts), batch_size):
-            batch_texts = new_texts[i : i+batch_size]
-            print(f"   Embedding batch {i//batch_size + 1}: Items {i} to {min(i+batch_size, len(new_texts))}...", flush=True)
-            emb_batch = model.encode(batch_texts, convert_to_numpy=True, normalize_embeddings=True)
-            new_embeddings_list.append(emb_batch)
-
-        new_embeddings = np.vstack(new_embeddings_list)
+            if chunk_new:
+                new_rows.extend(chunk_new)
+                new_cnt += len(chunk_new)
+            
+            scanned += len(df_chunk)
+            print(f"Scanned {scanned} rows... Found {new_cnt} new.", flush=True)
 
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Appending to database...", flush=True)
-        schema = pa.schema([('data', pa.string())])
-        new_table = pa.Table.from_pylist(new_rows, schema=schema)
+        logger.info(f"Total New Rows: {new_cnt}")
         
-        if first_run or not os.path.exists(PARQUET_FILE):
-            pq.write_table(new_table, PARQUET_FILE)
-            np.save(EMBEDDINGS_FILE, new_embeddings)
-        else:
-            old_table = pq.read_table(PARQUET_FILE)
-            combined_table = pa.concat_tables([old_table, new_table])
-            pq.write_table(combined_table, PARQUET_FILE)
+        # Save New Rows CSV
+        if new_rows:
+            pd.DataFrame(new_rows).to_csv(self.new_rows_csv, index=False)
             
-            old_embeddings = np.load(EMBEDDINGS_FILE)
-            combined_embeddings = np.vstack([old_embeddings, new_embeddings])
-            np.save(EMBEDDINGS_FILE, combined_embeddings)
+            # Append to Parquet
+            schema = pa.schema([('data', pa.string())])
+            new_table = pa.Table.from_pylist(new_rows, schema=schema)
             
+            if os.path.exists(self.local_parquet):
+                old_table = pq.read_table(self.local_parquet)
+                combined_table = pa.concat_tables([old_table, new_table])
+                pq.write_table(combined_table, self.full_updated_parquet)
+            else:
+                pq.write_table(new_table, self.full_updated_parquet)
+                
+        return new_cnt
+    
+    def cleanup_temp_files(self):
         if os.path.exists(self.temp_new_csv): os.remove(self.temp_new_csv)
-        print(f"[{datetime.now().strftime('%H:%M:%S')}] Job Complete.", flush=True)
-        return True
 
+# ==========================================
+# 3. EXECUTION BLOCK (Runs on Kaggle)
+# ==========================================
 if __name__ == "__main__":
+    
+    # A. Download Existing Data from HF (So your logic has something to compare)
+    REPO_ID = "ravi2814/grant-data-storage"
+    print("Downloading existing parquet from Hugging Face...", flush=True)
+    try:
+        hf_hub_download(repo_id=REPO_ID, filename="grant data.parquet", repo_type="dataset", local_dir=".", force_download=True)
+        hf_hub_download(repo_id=REPO_ID, filename="embeddings.npy", repo_type="dataset", local_dir=".", force_download=True)
+    except:
+        print("No existing data found. Starting fresh.", flush=True)
+
+    # B. Run Your Logic
     updater = GrantsDataUpdater()
-    updater.process()
+    
+    if updater.download_new_data():
+        hashes = updater.get_existing_hashes()
+        new_count = updater.process_and_find_new_rows(hashes)
+        
+        # C. Generate Embeddings for NEW rows (If any)
+        if new_count > 0:
+            print(f"Generating embeddings for {new_count} new rows...", flush=True)
+            device = 'cuda' if hasattr(sys, 'getandroidapilevel') else ('cuda' if os.path.exists('/kaggle') else 'cpu')
+            model = SentenceTransformer('intfloat/multilingual-e5-large', device=device)
+            
+            # Load new rows
+            df_new = pd.read_csv(updater.new_rows_csv)
+            texts = ["passage: " + str(t) for t in df_new['data']]
+            
+            new_embeddings = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=True)
+            
+            # Append Embeddings
+            if os.path.exists("embeddings.npy"):
+                old_embeddings = np.load("embeddings.npy")
+                combined = np.vstack([old_embeddings, new_embeddings])
+                np.save("embeddings.npy", combined)
+            else:
+                np.save("embeddings.npy", new_embeddings)
+            
+            print("Embeddings updated.", flush=True)
+        else:
+            print("No new rows to embed.", flush=True)
+            
+        updater.cleanup_temp_files()
+        print("Job Complete.", flush=True)