# loader.py # Load labeled data from Firebase storage for RLHF training import os import json import logging import pandas as pd from datetime import datetime from typing import List, Dict, Optional, Tuple, Any from pathlib import Path from io import StringIO # Import Firebase client from the existing firebase_saver import sys sys.path.append(os.path.join(os.path.dirname(__file__), '..')) from data.firebase_saver import _AdminClient, _GCSClient logger = logging.getLogger("rlhf-loader") logger.setLevel(logging.INFO) if not logger.handlers: _h = logging.StreamHandler() _h.setFormatter(logging.Formatter("[%(levelname)s] %(asctime)s - %(message)s")) logger.addHandler(_h) # Firebase configuration FIREBASE_BUCKET = "skyledge-36b56.firebasestorage.app" LABELED_PREFIX = "skyledge/labeled" RAW_PREFIX = "skyledge/raw" PROCESSED_PREFIX = "skyledge/processed" TRAINED_FILE = "trained.txt" class LabeledDataLoader: """ Load labeled data from Firebase storage for RLHF training. Tracks already processed datasets to avoid retraining on the same data. """ def __init__(self): self.bucket_name = FIREBASE_BUCKET self.prefix = LABELED_PREFIX self.trained_file = TRAINED_FILE # Initialize Firebase client self.client = None self.mode = None try: if os.getenv("FIREBASE_ADMIN_JSON"): self.client = _AdminClient(self.bucket_name) self.mode = "admin" except Exception as e: logger.warning(f"⚠️ Admin SDK init failed: {e}") if self.client is None: try: self.client = _GCSClient(self.bucket_name) self.mode = "gcs" except Exception as e: logger.error(f"❌ GCS client init failed: {e}") raise logger.info(f"📦 LabeledDataLoader ready | mode={self.mode} bucket={self.bucket_name} prefix={self.prefix}") def _get_trained_datasets(self) -> List[str]: """Load list of already trained datasets from trained.txt""" try: # Check if trained.txt exists in Firebase storage trained_path = f"{self.prefix}/{self.trained_file}" if self.client.blob_exists(trained_path): # Download and read the file blob = self.client.bucket.blob(trained_path) content = blob.download_as_text() trained_datasets = [line.strip() for line in content.split('\n') if line.strip()] logger.info(f"📋 Loaded {len(trained_datasets)} already trained datasets") return trained_datasets else: logger.info("📋 No trained.txt found, starting fresh") return [] except Exception as e: logger.warning(f"⚠️ Failed to load trained datasets: {e}") return [] def _update_trained_datasets(self, new_datasets: List[str]): """Update trained.txt with new dataset names""" try: # Get existing trained datasets existing = self._get_trained_datasets() # Add new datasets with timestamp timestamp = datetime.now().isoformat() new_entries = [f"{timestamp}:{dataset}" for dataset in new_datasets] all_entries = existing + new_entries # Upload updated file trained_path = f"{self.prefix}/{self.trained_file}" content = '\n'.join(all_entries) self.client.upload_from_bytes( content.encode('utf-8'), trained_path, "text/plain" ) logger.info(f"✅ Updated trained.txt with {len(new_datasets)} new datasets") except Exception as e: logger.error(f"❌ Failed to update trained datasets: {e}") def list_labeled_datasets(self) -> List[Dict[str, str]]: """List all available labeled datasets in Firebase storage""" try: # List all blobs under the labeled prefix blobs = list(self.client.bucket.list_blobs(prefix=f"{self.prefix}/")) datasets = [] trained_datasets = self._get_trained_datasets() for blob in blobs: # Skip the trained.txt file itself if blob.name.endswith(f"/{self.trained_file}"): continue # Extract dataset name (relative to skyledge root) dataset_name = blob.name.replace("skyledge/", "") # Skip if already trained if any(dataset_name in entry for entry in trained_datasets): continue # Get blob metadata blob.reload() datasets.append({ 'name': dataset_name, 'path': blob.name, 'size': blob.size, 'created': blob.time_created.isoformat() if blob.time_created else None, 'updated': blob.updated.isoformat() if blob.updated else None, 'content_type': blob.content_type }) logger.info(f"📊 Found {len(datasets)} new labeled datasets") return datasets except Exception as e: logger.error(f"❌ Failed to list labeled datasets: {e}") return [] def download_dataset(self, dataset_path: str, local_path: str) -> bool: """Download a dataset from Firebase storage to local path""" try: blob = self.client.bucket.blob(dataset_path) blob.download_to_filename(local_path) logger.info(f"✅ Downloaded {dataset_path} to {local_path}") return True except Exception as e: logger.error(f"❌ Failed to download {dataset_path}: {e}") return False def load_dataset(self, dataset_path: str) -> Optional[pd.DataFrame]: """Load a dataset directly into a pandas DataFrame""" try: blob = self.client.bucket.blob(dataset_path) content = blob.download_as_text() # Try to determine file type and load accordingly if dataset_path.endswith('.csv'): df = pd.read_csv(StringIO(content)) elif dataset_path.endswith('.json'): df = pd.read_json(StringIO(content)) elif dataset_path.endswith('.parquet'): # For parquet, we need to download as bytes blob_bytes = blob.download_as_bytes() df = pd.read_parquet(pd.BytesIO(blob_bytes)) else: # Default to CSV df = pd.read_csv(StringIO(content)) logger.info(f"✅ Loaded dataset {dataset_path} with shape {df.shape}") return df except Exception as e: logger.error(f"❌ Failed to load dataset {dataset_path}: {e}") return None def get_new_datasets_for_training(self) -> List[Dict[str, str]]: """Get list of new datasets that haven't been used for training yet""" return self.list_labeled_datasets() def mark_datasets_as_trained(self, dataset_names: List[str]): """Mark datasets as trained to avoid retraining""" self._update_trained_datasets(dataset_names) def _parse_labeled_filename(self, filename: str) -> Dict[str, str]: """ Parse labeled filename to extract original dataset information. Format: {id}_{source}-{original_id}_{date}-labelled.csv Example: 001_raw-002_2025-09-19-labelled.csv """ try: # Remove .csv extension name = filename.replace('.csv', '') # Split by underscore to get parts parts = name.split('_') if len(parts) < 4: return {"error": f"Invalid filename format: {filename}"} # Extract components labeled_id = parts[0] # 001 source_and_original = parts[1] # raw-002 or processed-002 date = parts[2] # 2025-09-19 # Parse source and original ID if '-' in source_and_original: source, original_id = source_and_original.split('-', 1) else: source = source_and_original original_id = "unknown" return { "labeled_id": labeled_id, "source": source, # raw or processed "original_id": original_id, "date": date, "original_filename": f"{original_id}_{date}-{source}.csv" if source != "unknown" else None } except Exception as e: logger.warning(f"⚠️ Failed to parse filename {filename}: {e}") return {"error": str(e)} def _find_original_dataset(self, labeled_info: Dict[str, str]) -> Optional[str]: """Find the original dataset path based on labeled file info""" if labeled_info.get("error") or not labeled_info.get("original_filename"): return None source = labeled_info["source"] original_filename = labeled_info["original_filename"] if source == "raw": return f"{self.RAW_PREFIX}/{original_filename}" elif source == "processed": return f"{self.PROCESSED_PREFIX}/{original_filename}" else: return None def load_labeled_with_original(self, labeled_path: str) -> Tuple[Optional[pd.DataFrame], Optional[pd.DataFrame], Dict[str, str]]: """ Load labeled dataset along with its original dataset for RLHF comparison. Returns: (labeled_df, original_df, metadata) """ try: # Load labeled dataset labeled_df = self.load_dataset(labeled_path) if labeled_df is None: return None, None, {"error": "Failed to load labeled dataset"} # Parse filename to get original dataset info filename = labeled_path.split('/')[-1] labeled_info = self._parse_labeled_filename(filename) if labeled_info.get("error"): logger.warning(f"⚠️ Could not parse labeled filename: {labeled_info['error']}") return labeled_df, None, labeled_info # Find and load original dataset original_path = self._find_original_dataset(labeled_info) original_df = None if original_path and self.client.blob_exists(original_path): original_df = self.load_dataset(original_path) if original_df is not None: logger.info(f"✅ Loaded original dataset: {original_path}") else: logger.warning(f"⚠️ Failed to load original dataset: {original_path}") else: logger.warning(f"⚠️ Original dataset not found: {original_path}") return labeled_df, original_df, labeled_info except Exception as e: logger.error(f"❌ Failed to load labeled with original: {e}") return None, None, {"error": str(e)} def create_training_batch(self, max_datasets: int = 10) -> Tuple[List[pd.DataFrame], List[str]]: """ Create a training batch by loading new datasets. Returns tuple of (dataframes, dataset_names) """ datasets = self.get_new_datasets_for_training() if not datasets: logger.info("📭 No new datasets available for training") return [], [] # Limit the number of datasets datasets = datasets[:max_datasets] dataframes = [] dataset_names = [] for dataset in datasets: df = self.load_dataset(dataset['path']) if df is not None: dataframes.append(df) dataset_names.append(dataset['name']) else: logger.warning(f"⚠️ Skipping dataset {dataset['name']} due to load failure") if dataframes: logger.info(f"📦 Created training batch with {len(dataframes)} datasets") # Mark these datasets as trained self.mark_datasets_as_trained(dataset_names) return dataframes, dataset_names def create_rlhf_training_batch(self, max_datasets: int = 10) -> Tuple[List[Dict[str, Any]], List[str]]: """ Create RLHF training batch with both labeled and original datasets. Returns tuple of (training_data, dataset_names) Each training_data item contains: {'labeled_df', 'original_df', 'metadata'} """ datasets = self.get_new_datasets_for_training() if not datasets: logger.info("📭 No new datasets available for RLHF training") return [], [] # Limit the number of datasets datasets = datasets[:max_datasets] training_data = [] dataset_names = [] for dataset in datasets: labeled_df, original_df, metadata = self.load_labeled_with_original(dataset['path']) if labeled_df is not None: training_item = { 'labeled_df': labeled_df, 'original_df': original_df, 'metadata': metadata, 'dataset_name': dataset['name'] } training_data.append(training_item) dataset_names.append(dataset['name']) logger.info(f"✅ Loaded RLHF dataset: {dataset['name']} (original: {metadata.get('original_filename', 'N/A')})") else: logger.warning(f"⚠️ Skipping dataset {dataset['name']} due to load failure") if training_data: logger.info(f"📦 Created RLHF training batch with {len(training_data)} datasets") # Mark these datasets as trained self.mark_datasets_as_trained(dataset_names) return training_data, dataset_names def main(): """Test the loader functionality""" loader = LabeledDataLoader() # List available datasets datasets = loader.list_labeled_datasets() print(f"Available datasets: {len(datasets)}") for dataset in datasets: print(f" - {dataset['name']} ({dataset['size']} bytes)") # Create a training batch dataframes, names = loader.create_training_batch(max_datasets=5) print(f"Training batch: {len(dataframes)} datasets") for name in names: print(f" - {name}") if __name__ == "__main__": main()