import os import sys import time import requests import json import shutil from huggingface_hub import HfApi, hf_hub_download from datetime import datetime # --- CONFIGURATION --- REPO_ID = "ravi2814/grant-data-storage" STATUS_FILE = "status.json" def get_token(): try: from kaggle_secrets import UserSecretsClient token = UserSecretsClient().get_secret("HF_TOKEN") if token: return token except: pass token = os.getenv("HF_TOKEN") if token and "PLACEHOLDER" not in token: return token return "HF_TOKEN_PLACEHOLDER" # --- STATUS MANAGER (Heartbeat) --- class StatusManager: def __init__(self, repo_id, token): self.repo_id = repo_id self.api = HfApi(token=token) self.last_upload = 0 def update(self, state, phase, logs, progress): """Uploads status.json to HF so the UI can read it remotely.""" try: # Rate limit: Update at most every 2 seconds if time.time() - self.last_upload < 2 and state == "running": return data = { "kaggle_state": state, "phase": phase, "logs": logs, "progress": progress, "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S") } # Write locally with open(STATUS_FILE, "w") as f: json.dump(data, f) # Upload to Dataset self.api.upload_file( path_or_fileobj=STATUS_FILE, path_in_repo=STATUS_FILE, repo_id=self.repo_id, repo_type="dataset" ) self.last_upload = time.time() except: pass # --- LOGGER --- class Logger: def __init__(self, repo_id, token, status_mgr): self.repo_id = repo_id self.api = HfApi(token=token) self.status_mgr = status_mgr self.log_file = "live_log.txt" self.buffer = [] self.last_upload = 0 # Init file with open(self.log_file, "w") as f: f.write(f"Worker Started at {datetime.now().strftime('%H:%M:%S')}\n") self.flush() def log(self, message): print(message, flush=True) timestamp = datetime.now().strftime('%H:%M:%S') log_entry = f"[{timestamp}] {message}" self.buffer.append(log_entry) # Update Status JSON immediately with log snippet self.status_mgr.update("running", "Connected to Kaggle (Safe to close)", log_entry, 95) if time.time() - self.last_upload > 3: self.flush() def flush(self): try: with open(self.log_file, "w") as f: f.write("\n".join(self.buffer)) self.api.upload_file(path_or_fileobj=self.log_file, path_in_repo=self.log_file, repo_id=self.repo_id, repo_type="dataset") self.last_upload = time.time() except: pass if __name__ == "__main__": HF_TOKEN = get_token() if "PLACEHOLDER" in HF_TOKEN: sys.exit(1) # 1. Start Status Manager status_mgr = StatusManager(REPO_ID, HF_TOKEN) status_mgr.update("running", "Initializing Environment...", "Starting...", 70) try: logger = Logger(REPO_ID, HF_TOKEN, status_mgr) logger.log("Environment Ready.") except: sys.exit(1) # 2. Install Libraries try: import sentence_transformers import pyarrow except ImportError: logger.log("Installing libraries...") os.system("pip install sentence-transformers pyarrow numpy pandas huggingface_hub torch") # 3. Imports (Fixed Torch Error) import pandas as pd import hashlib import re import numpy as np import pyarrow.parquet as pq import pyarrow as pa # Explicitly import torch with error handling try: import torch except ImportError: logger.log("Re-installing torch...") os.system("pip install torch") import torch from sentence_transformers import SentenceTransformer class Updater: def __init__(self, logger_instance): self.logger = logger_instance self.data_url = "https://open.canada.ca/data/dataset/432527ab-7aac-45b5-81d6-7597107a7013/resource/1d15a62f-5656-49ad-8c88-f40ce689d831/download/grants.csv" self.api_url = "https://open.canada.ca/data/api/3/action/package_show?id=432527ab-7aac-45b5-81d6-7597107a7013" self.local_parquet = "grant data.parquet" self.temp_csv = "temp_data.csv" self.api_metadata = {} device = 'cuda' if torch.cuda.is_available() else 'cpu' self.logger.log(f"Loading AI Model on {device.upper()}...") self.model = SentenceTransformer('intfloat/multilingual-e5-large', device=device) def create_hash(self, text): if not isinstance(text, str): text = "" clean = re.sub(r'\s+', '', text).lower() return hashlib.md5(clean.encode('utf-8')).hexdigest() def format_record(self, record): def val(k): return str(record.get(k, "") or "").strip() mapping = [ ("Reference number", "ref_number"), ("amendment number", "amendment_number"), ("agreement type", "agreement_type"), ("recipient legal name", "recipient_legal_name"), ("recipient country", "recipient_country"), ("recipient province", "recipient_province"), ("recipient city", "recipient_city"), ("recipient postal code", "recipient_postal_code"), ("program name english", "prog_name_en"), ("program name french", "prog_name_fr"), ("program purpose english", "prog_purpose_en"), ("program purpose french", "prog_purpose_fr"), ("agreement title english", "agreement_title_en"), ("agreement title french", "agreement_title_fr"), ("agreement number", "agreement_number"), ("agreement value", "agreement_value"), ("agreement start date", "agreement_start_date"), ("agreement end date", "agreement_end_date"), ("description english", "description_en"), ("description french", "description_fr"), ("expected results english", "expected_results_en"), ("expected results french", "expected_results_fr"), ("additional information english", "additional_information_en"), ("additional information french", "additional_information_fr"), ("owner organization", "owner_org"), ("owner organization title", "owner_org_title"), ] return ", ".join([f"{d} : {val(k)}" for d, k in mapping]) def download_current_database(self): self.logger.log("Downloading current Database...") try: hf_hub_download(repo_id=REPO_ID, filename="grant data.parquet", repo_type="dataset", token=HF_TOKEN, local_dir=".", force_download=True) hf_hub_download(repo_id=REPO_ID, filename="embeddings.npy", repo_type="dataset", token=HF_TOKEN, local_dir=".", force_download=True) except: self.logger.log("No existing database found. Starting fresh.") def download_government_data(self): self.logger.log("Downloading fresh data from Canada.ca...") try: r = requests.get(self.data_url, stream=True) with open(self.temp_csv, 'wb') as f: for chunk in r.iter_content(chunk_size=1048576): f.write(chunk) r_api = requests.get(self.api_url).json()['result'] res = next((r for r in r_api['resources'] if r['id'] == '1d15a62f-5656-49ad-8c88-f40ce689d831'), {}) self.api_metadata = { "metadata_modified": r_api.get('metadata_modified'), "resource_modified": res.get('last_modified'), "checked_at": datetime.now().isoformat() } return True except Exception as e: self.logger.log(f"Download Error: {e}") return False def synchronize_data(self): if not os.path.exists(self.temp_csv): return False self.logger.log("Scanning Government CSV...") status_mgr.update("running", "Scanning CSV...", "Reading...", 60) valid_hashes = set() new_rows_buffer = [] for chunk in pd.read_csv(self.temp_csv, chunksize=50000, low_memory=False): for _, row in chunk.iterrows(): fmt = self.format_record(row.to_dict()) h = self.create_hash(fmt) valid_hashes.add(h) new_rows_buffer.append({'data': fmt, 'hash': h}) self.logger.log(f"Total Official Records: {len(valid_hashes)}") if os.path.exists(self.local_parquet) and os.path.exists("embeddings.npy"): df = pd.read_parquet(self.local_parquet) emb = np.load("embeddings.npy") keep_indices = [] local_hashes = set() for idx, row in df.iterrows(): h = self.create_hash(row['data']) if h in valid_hashes: keep_indices.append(idx) local_hashes.add(h) if len(keep_indices) < len(df): self.logger.log(f"Removing {len(df) - len(keep_indices)} obsolete records...") df = df.loc[keep_indices].reset_index(drop=True) emb = emb[keep_indices] else: df = pd.DataFrame(columns=['data']) emb = np.empty((0, 1024)) local_hashes = set() rows_to_add = [] for item in new_rows_buffer: if item['hash'] not in local_hashes: rows_to_add.append(item['data']) local_hashes.add(item['hash']) if rows_to_add: self.logger.log(f"Found {len(rows_to_add)} NEW records. Processing...") status_mgr.update("running", "Embedding New Data...", f"Processing {len(rows_to_add)} rows", 80) new_vectors = [] for i in range(0, len(rows_to_add), 2000): batch = rows_to_add[i : i + 2000] batch_fmt = ["passage: " + t for t in batch] v = self.model.encode(batch_fmt, convert_to_numpy=True, normalize_embeddings=True) new_vectors.append(v) new_emb_stack = np.vstack(new_vectors) new_df = pd.DataFrame({'data': rows_to_add}) df = pd.concat([df, new_df], ignore_index=True) if len(emb) > 0: emb = np.vstack([emb, new_emb_stack]) else: emb = new_emb_stack self.logger.log(f"Database Updated. Total: {len(df)}") else: self.logger.log("No new records to add.") self.logger.log("Saving files...") table = pa.Table.from_pandas(df) pq.write_table(table, self.local_parquet) np.save("embeddings.npy", emb) with open("last_metadata.json", "w") as f: json.dump(self.api_metadata, f) return True try: updater = Updater(logger) updater.download_current_database() if updater.download_government_data(): updater.synchronize_data() logger.log("Uploading results...") status_mgr.update("running", "Uploading Final Data...", "Finalizing...", 95) api = HfApi(token=HF_TOKEN) os.makedirs("staging", exist_ok=True) files = ["grant data.parquet", "embeddings.npy", "last_metadata.json", "live_log.txt", STATUS_FILE] for f in files: if os.path.exists(f): shutil.copy(f, f"staging/{f}") api.upload_folder(folder_path="staging", repo_id=REPO_ID, repo_type="dataset", commit_message="Batch Update") # FINAL SUCCESS status_mgr.update("finished", "Sync Complete", "All data uploaded successfully.", 100) logger.log("SUCCESS: Process Complete.") except Exception as e: logger.log(f"CRITICAL ERROR: {e}") status_mgr.update("error", "Worker Failed", str(e), 0) sys.exit(1)