Spaces:
Running on Zero
Running on Zero
| """Layer metadata and neuron detail loading with on-demand caching.""" | |
| import json | |
| import os | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from datasets import get_dataset_split_names, load_dataset | |
| from config import QWEN25_DATASET_ROOT, REPO_ID, qwen25_dataset_path | |
| # Config names and directories in the HF dataset | |
| CONFIG_META = f"{QWEN25_DATASET_ROOT}_meta_data" | |
| CONFIG_DETAILS = f"{QWEN25_DATASET_ROOT}_neuron_details" | |
| # Global cache for per-layer meta dict (lightweight) | |
| layer_meta_cache: Dict[int, Dict[int, Dict]] = {} | |
| # Cache for per-layer, per-chunk neuron detail datasets (non-streaming) | |
| # Keyed by layer_id -> chunk_start (every 1000 neurons) -> dataset split object | |
| layer_datasets: Dict[int, Dict[int, Any]] = {} | |
| # Cache for per-neuron detailed rows (only selected neuron) | |
| neuron_detail_cache: Dict[Tuple[int, int], Dict] = {} | |
| def _summarize_splits(split_names: List[str], head: int = 5, tail: int = 5) -> str: | |
| if not split_names: | |
| return "[]" | |
| if len(split_names) <= head + tail: | |
| return str(split_names) | |
| return str(split_names[:head] + ["..."] + split_names[-tail:]) | |
| def load_layer_options() -> List[int]: | |
| """Fetch available layer IDs, prefer local JSON, fallback to default.""" | |
| try: | |
| if os.path.exists("resources/layer_list.json"): | |
| with open("resources/layer_list.json", "r", encoding="utf-8") as f: | |
| layer_options = json.load(f) | |
| layer_options = sorted(int(x) for x in layer_options) | |
| else: | |
| layer_options = [20] | |
| print(f"[info] Available layer IDs: {layer_options}") | |
| return layer_options | |
| except Exception as e: # pragma: no cover - small utility | |
| print(f"[error] Failed to load layer options: {e}") | |
| return [20] | |
| def load_layer_data(layer_id: int) -> Dict[int, Dict]: | |
| """Load lightweight meta data for a specific layer from HF dataset. | |
| Uses a single-file data_files load for layer meta parquet. | |
| Only minimal fields are stored; heavy refs are not loaded here. | |
| """ | |
| if layer_id in layer_meta_cache: | |
| cached = layer_meta_cache[layer_id] | |
| print(f"[info] Using cached meta_data for layer {layer_id} (size={len(cached)})") | |
| return cached | |
| try: | |
| data_file = qwen25_dataset_path( | |
| f"meta_data/layer_{layer_id}-00000-of-00001.parquet" | |
| ) | |
| print(f"[info] Loading meta_data via data_files: layer={layer_id} file={data_file}") | |
| ds = load_dataset( | |
| REPO_ID, | |
| data_files={"train": [data_file]}, | |
| split="train", | |
| streaming=False, | |
| ) | |
| print(f"[info] meta_data loaded: file={data_file} rows={len(ds)}") | |
| meta_fields = { | |
| "neuron_id", | |
| "modality", | |
| "vision_summary", | |
| "language_summary", | |
| "consistency_score", | |
| "miou", | |
| } | |
| meta_dict: Dict[int, Dict] = {} | |
| for item in ds: | |
| nid = item.get("neuron_id") | |
| if nid is None: | |
| continue | |
| try: | |
| meta = {k: item.get(k) for k in meta_fields} | |
| meta_dict[int(nid)] = meta | |
| except Exception as row_err: # pragma: no cover | |
| print(f"[warn] Skipping meta row for layer {layer_id}: {row_err}") | |
| layer_meta_cache[layer_id] = meta_dict | |
| print(f"[info] Loaded meta_data for {len(meta_dict)} neurons for layer {layer_id}") | |
| return meta_dict | |
| except Exception as e: # pragma: no cover - network/cache dependent | |
| print(f"[error] Failed to load meta_data for layer {layer_id}: {e}") | |
| return {} | |
| def _details_split_name(layer_id: int, neuron_id: int, chunk: int = 1000) -> str: | |
| start = (int(neuron_id) // chunk) * chunk | |
| end = start + (chunk - 1) | |
| return f"layer_{layer_id}_neuron_{start:06d}_{end:06d}" | |
| def load_neuron_detail(layer_id: int, neuron_id: int) -> Optional[Dict]: | |
| """Load detailed refs (vision/text) for the selected neuron only. | |
| The dataset index now aligns with the neuron index, including empty neurons, | |
| so we can directly index into the split without an extra scan on `neuron_id`. | |
| """ | |
| cache_key = (int(layer_id), int(neuron_id)) | |
| if cache_key in neuron_detail_cache: | |
| print(f"[info] Using cached neuron_details for neuron {neuron_id} layer {layer_id}") | |
| return neuron_detail_cache[cache_key] | |
| # Ensure meta exists (for sanity), but don't depend on its split info | |
| meta_dict = load_layer_data(int(layer_id)) | |
| if int(neuron_id) not in meta_dict: | |
| return None | |
| chunk_size = 1000 | |
| split_name = _details_split_name(int(layer_id), int(neuron_id), chunk=chunk_size) | |
| print(f"[info] Loading neuron_details: layer={layer_id} neuron={neuron_id} split={split_name}") | |
| split_names = None | |
| try: | |
| split_names = get_dataset_split_names(path=REPO_ID, config_name=CONFIG_DETAILS) | |
| print( | |
| f"[info] neuron_details splits via config='{CONFIG_DETAILS}': count={len(split_names)} " | |
| f"sample={_summarize_splits(split_names)}" | |
| ) | |
| except Exception as e: | |
| print( | |
| f"[warn] Failed to list neuron_details splits via config: {e}; " | |
| "trying the expected file directly..." | |
| ) | |
| if split_names is not None and split_name not in split_names: | |
| print(f"[warn] neuron_details split '{split_name}' not found for layer {layer_id}") | |
| return None | |
| try: | |
| chunk_start = (int(neuron_id) // chunk_size) * chunk_size | |
| layer_cache = layer_datasets.setdefault(int(layer_id), {}) | |
| if chunk_start in layer_cache: | |
| ds = layer_cache[chunk_start] | |
| print( | |
| f"[info] Using cached neuron_details split='{split_name}' for layer {layer_id} " | |
| f"chunk_start={chunk_start}" | |
| ) | |
| else: | |
| data_file = qwen25_dataset_path( | |
| f"neuron_details/layer_{layer_id}/{split_name}-00000-of-00001.parquet" | |
| ) | |
| print( | |
| f"[info] Loading neuron_details via data_files: layer={layer_id} split={split_name} " | |
| f"file={data_file}" | |
| ) | |
| ds = load_dataset( | |
| REPO_ID, | |
| data_files={"train": [data_file]}, | |
| split="train", | |
| streaming=False, | |
| ) | |
| layer_cache[chunk_start] = ds | |
| print(f"[info] Cached split='{split_name}' rows={len(ds)} for layer {layer_id}") | |
| # Dataset index matches neuron index; empty neurons yield empty rows. | |
| local_index = int(neuron_id) - chunk_start | |
| if local_index < 0 or local_index >= len(ds): | |
| print( | |
| f"[warn] Neuron {neuron_id} maps to local index {local_index}, " | |
| f"but split '{split_name}' has length {len(ds)}" | |
| ) | |
| return None | |
| row = ds[int(local_index)] | |
| neuron_detail_cache[cache_key] = row | |
| print(f"[info] Loaded neuron_details for neuron {neuron_id} layer {layer_id} via direct index") | |
| return row | |
| except Exception as e: # pragma: no cover - network/cache dependent | |
| print(f"[error] Failed to load neuron_details for neuron {neuron_id} layer {layer_id}: {e}") | |
| return None | |
| def get_neuron_ids(layer_id: int) -> List[int]: | |
| """Return all neuron IDs for the given layer from meta_data.""" | |
| data = load_layer_data(layer_id) | |
| if data: | |
| return sorted(data.keys()) | |
| return [] | |
| def get_consistency_neurons(layer_id: int, max_examples: int = 25) -> List[int]: | |
| """Return up to max_examples neurons with consistency_score == 1 from meta_data.""" | |
| data = load_layer_data(layer_id) | |
| if not data: | |
| return [] | |
| matches = [nid for nid, meta in data.items() if meta.get("consistency_score") == 1.0] | |
| return sorted(matches)[:max_examples] | |