Upload folder using huggingface_hub

data/data_collator.py

@@ -711,11 +711,15 @@ class MemecoinCollator:
             # Labels
             'labels': torch.stack([item['labels'] for item in batch]) if batch and 'labels' in batch[0] else None,
             'labels_mask': torch.stack([item['labels_mask'] for item in batch]) if batch and 'labels_mask' in batch[0] else None,
+            'quality_score': torch.stack([item['quality_score'] for item in batch]) if batch and 'quality_score' in batch[0] else None,
             # Debug info
             'token_addresses': [item.get('token_address', 'unknown') for item in batch],
             't_cutoffs': [item.get('t_cutoff', 'unknown') for item in batch],
             'sample_indices': [item.get('sample_idx', -1) for item in batch]
         }
 
+        if collated_batch['quality_score'] is None:
+            raise RuntimeError("FATAL: Missing quality_score in batch items. Rebuild cache with quality_score enabled.")
+
         # Filter out None values (e.g., if no labels provided)
         return {k: v for k, v in collated_batch.items() if v is not None}

data/data_loader.py

@@ -156,43 +156,41 @@ class OracleDataset(Dataset):
             if not self.cached_files:
                 raise RuntimeError(f"Cache directory '{self.cache_dir}' provided but contains no 'sample_*.pt' files.")
 
-            # --- NEW: Strict Metadata & Weighting ---
-            metadata_path = self.cache_dir / "metadata.jsonl"
-            if not metadata_path.exists():
-                raise RuntimeError(f"FATAL: metadata.jsonl not found in {self.cache_dir}. Cannot train without class-balanced sampling.")
-            
-            print(f"INFO: Loading metadata from {metadata_path}...")
+            # --- NEW: Strict Metadata & Weighting (from cached samples) ---
             file_class_map = {}
             class_counts = defaultdict(int)
-            
-            with open(metadata_path, 'r') as f:
-                for line in f:
+
+            # Read class_id directly from each cached sample
+            for p in self.cached_files:
+                try:
+                    # Cached samples are trusted local artifacts; allow full load.
                     try:
-                        entry = json.loads(line)
-                        fname = entry['file']
-                        cid = entry['class_id']
-                        file_class_map[fname] = cid
-                        class_counts[cid] += 1
-                    except Exception as e:
-                        print(f"WARN: Failed to parse metadata line: {e}")
+                        cached_item = torch.load(p, map_location="cpu", weights_only=False)
+                    except TypeError:
+                        cached_item = torch.load(p, map_location="cpu")
+                    cid = cached_item.get("class_id")
+                    if cid is None:
+                        print(f"WARN: File {p.name} missing class_id. Skipping.")
+                        continue
+                    file_class_map[p.name] = cid
+                    class_counts[cid] += 1
+                except Exception as e:
+                    print(f"WARN: Failed to read cached sample {p.name}: {e}")
 
             print(f"INFO: Class Distribution: {dict(class_counts)}")
-            
+
             # Compute Weights
             self.weights_list = []
             valid_files = []
-            
+
             # We iterate properly sorted cached files to align with __getitem__ index
             for p in self.cached_files:
                 fname = p.name
                 if fname not in file_class_map:
-                    # Should be fatal if strict, but maybe some files were skipped?
-                    # If file exists but no metadata, we can't weight it properly.
-                    # Current pipeline writes metadata only for successful caches.
-                    # So if it's in cached_files but not metadata, it might be a stale file.
-                    print(f"WARN: File {fname} found in cache but missing metadata. Skipping.")
+                    # If file exists but missing class_id, it might be stale or from an older cache.
+                    print(f"WARN: File {fname} found in cache but missing class_id. Skipping.")
                     continue
-                
+
                 cid = file_class_map[fname]
                 count = class_counts[cid]
                 weight = 1.0 / count if count > 0 else 0.0
@@ -976,7 +974,8 @@ class OracleDataset(Dataset):
             "fee_collections",
             "burns",
             "supply_locks",
-            "migrations"
+            "migrations",
+            "quality_score"
         ]
         missing_keys = [key for key in required_keys if key not in raw_data]
         if missing_keys:
@@ -1683,7 +1682,8 @@ class OracleDataset(Dataset):
                 'graph_links': graph_links,
                 'embedding_pooler': pooler,
                 'labels': torch.zeros(len(self.horizons_seconds), dtype=torch.float32),
-                'labels_mask': torch.zeros(len(self.horizons_seconds), dtype=torch.float32)
+                'labels_mask': torch.zeros(len(self.horizons_seconds), dtype=torch.float32),
+                'quality_score': torch.tensor(raw_data['quality_score'], dtype=torch.float32)
             }
         
         # Ensure sorted
@@ -1758,5 +1758,6 @@ class OracleDataset(Dataset):
             'graph_links': graph_links, 
             'embedding_pooler': pooler, 
             'labels': torch.tensor(label_values, dtype=torch.float32),
-            'labels_mask': torch.tensor(mask_values, dtype=torch.float32)
+            'labels_mask': torch.tensor(mask_values, dtype=torch.float32),
+            'quality_score': torch.tensor(raw_data['quality_score'], dtype=torch.float32)
         }

data/ohlc_stats.npz

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e8366dfe6785692219a4d4bcbe5c3b111b5b9acd3df38fba7edd5d29bea20e15
+oid sha256:6f2c86bf03e5761e7fb319a54274e032f7aa1d01dd5873f2f44a52c9e0be5244
 size 1660

log.log

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b4dd7b51859975e9b53550cdda3099bd1fd899d8b335ff3b90ab5ae7d9a1e4ff
-size 4414
+oid sha256:461e55d31752fd72f09aa30c5bcc3a619654ae86ddf1e759c9c57b0dc5db53f6
+size 21794

models/model.py

@@ -54,7 +54,9 @@ class Oracle(nn.Module):
         self.dtype = dtype
         
         # --- 2. Load Qwen3 Configuration (architecture only; training from scratch) ---
-        model_config = AutoConfig.from_pretrained(model_config_name, trust_remote_code=True)
+        hf_token = os.getenv("Hf_TOKEN") or os.getenv("HF_TOKEN")
+        hf_kwargs = {"token": hf_token} if hf_token else {}
+        model_config = AutoConfig.from_pretrained(model_config_name, trust_remote_code=True, **hf_kwargs)
         self.d_model = model_config.hidden_size
         self.model = AutoModel.from_config(model_config, trust_remote_code=True)
         self.model.to(self.device, dtype=self.dtype)
@@ -65,6 +67,11 @@ class Oracle(nn.Module):
             nn.GELU(),
             nn.Linear(self.d_model, self.num_outputs)
         )
+        self.quality_head = nn.Sequential(
+            nn.Linear(self.d_model, self.d_model),
+            nn.GELU(),
+            nn.Linear(self.d_model, 1)
+        )
 
         self.event_type_to_id = event_type_to_id
 
@@ -947,8 +954,10 @@ class Oracle(nn.Module):
             empty_hidden = torch.empty(0, L, self.d_model, device=device, dtype=self.dtype)
             empty_mask = torch.empty(0, L, device=device, dtype=torch.long)
             empty_quantiles = torch.empty(0, self.num_outputs, device=device, dtype=self.dtype)
+            empty_quality = torch.empty(0, device=device, dtype=self.dtype)
             return {
                 'quantile_logits': empty_quantiles,
+                'quality_logits': empty_quality,
                 'pooled_states': torch.empty(0, self.d_model, device=device, dtype=self.dtype),
                 'hidden_states': empty_hidden,
                 'attention_mask': empty_mask
@@ -1068,9 +1077,11 @@ class Oracle(nn.Module):
         sequence_hidden = outputs.last_hidden_state
         pooled_states = self._pool_hidden_states(sequence_hidden, hf_attention_mask)
         quantile_logits = self.quantile_head(pooled_states)
+        quality_logits = self.quality_head(pooled_states).squeeze(-1)
 
         return {
             'quantile_logits': quantile_logits,
+            'quality_logits': quality_logits,
             'pooled_states': pooled_states,
             'hidden_states': sequence_hidden,
             'attention_mask': hf_attention_mask

models/multi_modal_processor.py

@@ -38,13 +38,16 @@ class MultiModalEncoder:
 
 
         try:
+            hf_token = os.getenv("Hf_TOKEN") or os.getenv("HF_TOKEN")
+            hf_kwargs = {"token": hf_token} if hf_token else {}
             # --- SigLIP Loading with Config Fix ---
             self.processor = AutoProcessor.from_pretrained(
                 self.model_id,
-                use_fast=True
+                use_fast=True,
+                **hf_kwargs
             )
 
-            config = AutoConfig.from_pretrained(self.model_id)
+            config = AutoConfig.from_pretrained(self.model_id, **hf_kwargs)
 
             if not hasattr(config, 'projection_dim'):
                 # print("❗ Config missing projection_dim, patching...")
@@ -54,7 +57,8 @@ class MultiModalEncoder:
                 self.model_id,
                 config=config,
                 dtype=self.dtype, # Use torch_dtype for from_pretrained
-                trust_remote_code=False
+                trust_remote_code=False,
+                **hf_kwargs
             ).to(self.device).eval()
             # -----------------------------------------------
 

pre_cache.sh

@@ -4,6 +4,6 @@
 echo "Starting dataset caching..."
 python3 scripts/cache_dataset.py \
     --ohlc_stats_path "/workspace/apollo/data/ohlc_stats.npz" \
-    --max_samples 1000
+    --max_samples 50
 
 echo "Done!"

scripts/analyze_distribution.py

@@ -1,21 +1,22 @@
-
 import os
 import sys
 import datetime
+import numpy as np
+import math
 from clickhouse_driver import Client as ClickHouseClient
 
 # Add parent to path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-# removed dotenv
-# load_dotenv()
+from models.vocabulary import RETURN_THRESHOLDS, MANIPULATED_CLASS_ID
 
 CLICKHOUSE_HOST = os.getenv("CLICKHOUSE_HOST", "localhost")
 CLICKHOUSE_PORT = int(os.getenv("CLICKHOUSE_PORT", 9000))
-# .env shows empty user/pass, which implies 'default' user and empty password for ClickHouse
 CLICKHOUSE_USER = os.getenv("CLICKHOUSE_USER", "default") 
 CLICKHOUSE_PASSWORD = os.getenv("CLICKHOUSE_PASSWORD", "")
 CLICKHOUSE_DATABASE = os.getenv("CLICKHOUSE_DATABASE", "default")
+LAUNCH_PRICE_USD = 0.000004
+EPS = 1e-9
 
 def get_client():
     return ClickHouseClient(
@@ -26,484 +27,331 @@ def get_client():
         database=CLICKHOUSE_DATABASE
     )
 
-def print_distribution_stats(client, metric_name, subquery, bucket_case_sql):
-    print(f"\n   -> {metric_name}")
-    
-    # 1. Print Basic Stats (Mean, Quantiles)
-    stats_query = f"""
-    SELECT
-        avg(val),
-        quantiles(0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99)(val),
-        min(val),
-        max(val),
-        count()
-    FROM (
-        {subquery}
-    )
+def fetch_all_metrics(client):
+    """
+    Fetches all needed metrics for all tokens in a single query.
+    Base Table: MINTS (to ensure we cover all ~50k tokens).
+    Definitions:
+      - Snipers: Peak Balance Sum of top 70 buyers
+      - Bundles: Base Amount Sum of trades in multi-buy slots
+      - Dev Hold: Max Peak Balance of Creator
     """
-    try:
-        stats = client.execute(stats_query)[0]
-        avg_val = stats[0]
-        qs = stats[1]
-        min_val = stats[2]
-        max_val = stats[3]
-        count_val = stats[4]
+    print("   -> Fetching all token metrics (Unified Query)...")
+    
+    query = f"""
+    WITH
+        -- 1. Aggregated trade stats (Fees, Volume, ATH Time)
+        trade_agg AS (
+            SELECT
+                base_address,
+                sum(priority_fee + coin_creator_fee) AS fees_sol,
+                sum(total_usd) AS volume_usd,
+                count() AS n_trades,
+                argMax(timestamp, price_usd) AS t_ath,
+                min(timestamp) AS t0
+            FROM trades
+            GROUP BY base_address
+        ),
+        
+        -- 2. Token Metadata from MINTS (Base Source of Truth)
+        token_meta AS (
+            SELECT
+                mint_address AS token_address,
+                argMax(creator_address, timestamp) AS creator_address,
+                argMax(total_supply, timestamp) AS total_supply,
+                argMax(token_decimals, timestamp) AS decimals
+            FROM mints
+            GROUP BY mint_address
+        ),
         
-        if count_val == 0:
-            print("      No data for this segment.")
-            return
+        -- 3. Returns & Holders (from Token Metrics or manual calc)
+        metrics AS (
+             SELECT
+                token_address,
+                argMax(ath_price_usd, updated_at) as ath_price_usd,
+                argMax(unique_holders, updated_at) as unique_holders
+             FROM token_metrics
+             GROUP BY token_address
+        ),
 
-        print(f"      Mean: {avg_val:.4f} | Min: {min_val:.4f} | Max: {max_val:.4f}")
-        print(f"      Q: p10={qs[0]:.2f} p50={qs[2]:.2f} p90={qs[4]:.2f} p99={qs[6]:.2f}")
+        -- 4. WALLET PEAKS (normalized balance likely)
+        wallet_peaks AS (
+            SELECT
+                mint_address,
+                wallet_address,
+                max(current_balance) AS peak_balance
+            FROM wallet_holdings
+            GROUP BY mint_address, wallet_address
+        ),
 
-    except Exception as e:
-        print(f"      Error calculating stats: {e}")
-        return
+        -- 5. SNIPERS: Identify sniper addresses (rank <= 70)
+        snipers_list AS (
+             SELECT
+                 base_address,
+                 maker
+             FROM (
+                 SELECT
+                     base_address,
+                     maker,
+                     dense_rank() OVER (PARTITION BY base_address ORDER BY min_slot, min_idx) AS buyer_rank
+                 FROM (
+                     SELECT
+                         base_address,
+                         maker,
+                         min(slot) AS min_slot,
+                         min(transaction_index) AS min_idx
+                     FROM trades
+                     WHERE trade_type = 0 -- buy
+                     GROUP BY base_address, maker
+                 )
+             )
+             WHERE buyer_rank <= 70
+        ),
+        snipers_agg AS (
+            SELECT
+                s.base_address AS token_address,
+                sum(wp.peak_balance) AS snipers_total_peak
+            FROM snipers_list s
+            JOIN wallet_peaks wp ON s.base_address = wp.mint_address AND s.maker = wp.wallet_address
+            GROUP BY s.base_address
+        ),
 
-    # 2. Print Buckets
-    query = f"""
-    SELECT
-        {bucket_case_sql} as bucket,
-        count() as cnt
-    FROM (
-        {subquery}
-    )
-    GROUP BY bucket
-    ORDER BY bucket
-    """
-    try:
-        rows = client.execute(query)
-        # total_count used for pct is the count_val from stats
-        print("      Buckets:")
-        for r in rows:
-            pct = (r[1] / count_val * 100) if count_val > 0 else 0
-            print(f"        {r[0]}: {r[1]} ({pct:.1f}%)")
-    except Exception as e:
-        print(f"      Error calculating buckets: {e}")
+        -- 6. BUNDLED: Sum the base_amount of ALL trades that happened in a slot with multiple buys
+        bundled_agg AS (
+            SELECT
+                t.base_address AS token_address,
+                sum(t.base_amount) AS bundled_total_peak
+            FROM trades t
+            WHERE (t.base_address, t.slot) IN (
+                 SELECT base_address, slot
+                 FROM trades
+                 WHERE trade_type = 0 -- buy
+                 GROUP BY base_address, slot
+                 HAVING count() > 1
+            )
+            AND t.trade_type = 0 -- buy
+            GROUP BY t.base_address
+        ),
 
-def get_filtered_metric_query(inner_query, cohort_sql):
-    """
-    Wraps the inner metric query to only include tokens in the cohort.
-    Assumes inner_query returns 'base_address' (or aliased) and 'val'.
-    If the inner query returns 'token_address', it should be handled.
-    Most of our queries return 'base_address' (from trades) or 'token_address' (from token_metrics).
-    We will normalize to use 'base_address' via subquery alias if needed, but simplest is 
-    to filter on the outer Select.
+        -- 7. DEV HOLD: Creator's Peak Balance
+        dev_hold_agg AS (
+            SELECT
+                t.token_address,
+                max(wp.peak_balance) AS dev_peak
+            FROM token_meta t
+            JOIN wallet_peaks wp ON t.token_address = wp.mint_address AND t.creator_address = wp.wallet_address
+            GROUP BY t.token_address
+        )
+
+    SELECT
+        t.token_address,
+        (COALESCE(m.ath_price_usd, ta.t_ath, 0) / {LAUNCH_PRICE_USD}) AS ret,
+        
+        COALESCE(ta.fees_sol, 0) AS fees_sol,
+        COALESCE(ta.volume_usd, 0) AS volume_usd,
+        COALESCE(m.unique_holders, 0) AS unique_holders,
+        (ta.t_ath - ta.t0) AS time_to_ath_sec,
+        
+        COALESCE(s.snipers_total_peak, 0) AS snipers_val,
+        COALESCE(b.bundled_total_peak, 0) AS bundled_val,
+        COALESCE(d.dev_peak, 0) AS dev_val,
+        
+        t.total_supply AS total_supply,
+        t.decimals AS decimals
+        
+    FROM token_meta t
+    LEFT JOIN trade_agg ta ON t.token_address = ta.base_address
+    LEFT JOIN metrics m ON t.token_address = m.token_address
+    LEFT JOIN snipers_agg s ON t.token_address = s.token_address
+    LEFT JOIN bundled_agg b ON t.token_address = b.token_address
+    LEFT JOIN dev_hold_agg d ON t.token_address = d.token_address
     """
-    # We need to know if the inner query produces 'base_address' or 'token_address'
-    # Currently our queries produce 'base_address' mostly, except token_metrics ones.
-    # Let's standardize inner queries in the main loop to alias the key column to 'join_key'
     
-    return f"""
-    SELECT * FROM (
-        {inner_query}
-    ) WHERE join_key IN ({cohort_sql})
-    """
+    rows = client.execute(query)
+    # Convert to list of dicts
+    cols = [
+        "token_address", "ret", "fees_sol", "volume_usd", "unique_holders", "time_to_ath_sec",
+        "snipers_val", "bundled_val", "dev_val", "total_supply", "decimals"
+    ]
+    results = []
+    
+    print(f"   -> Fetched {len(rows)} tokens.")
+    
+    for r in rows:
+        d = dict(zip(cols, r))
+        
+        supply = d["total_supply"]
+        decimals = d["decimals"]
+        
+        try:
+            adj_supply = supply / (10 ** decimals) if (supply and decimals is not None) else supply
+        except:
+            adj_supply = supply
 
-import numpy as np
-from models.vocabulary import RETURN_THRESHOLDS, MANIPULATED_CLASS_ID
+        if adj_supply and adj_supply > 0:
+            d["snipers_pct"] = (d["snipers_val"] / adj_supply) * 100
+            d["dev_hold_pct"] = (d["dev_val"] / adj_supply) * 100
+        else:
+            d["snipers_pct"] = 0.0
+            d["dev_hold_pct"] = 0.0
+            
+        if supply and supply > 0:
+            d["bundled_pct"] = (d["bundled_val"] / supply) * 100
+        else:
+            d["bundled_pct"] = 0.0
 
-def get_return_class_map(client):
-    """
-    Returns a dictionary mapping token_address -> class_id (int)
-    Filters out tokens with > 10,000x return.
-    Implements Dynamic Outlier Detection:
-    - Calculates Median Fees, Volume, Holders for each Class (1-4).
-    - Downgrades tokens with metrics < 10% of their class median to Class 5 (Manipulated).
+        results.append(d)
+        
+    return results
+
+def _classify_tokens(data):
     """
-    print("   -> Fetching metrics for classification...")
-    # improved query to get fees/vol/holders
-    # aggregating trades for fees/vol to appear more robust than token_metrics snapshots
-    print("   -> Fetching metrics for classification...")
-    # SQL OPTIMIZATION:
-    # 1. Use token_metrics for Volume/Holders (Pre-computed).
-    # 2. Pre-aggregate trades for Fees in a subquery to avoid massive JOIN explosion.
-    query = """
-    SELECT 
-        tm.token_address, 
-        (argMax(tm.ath_price_usd, tm.updated_at) / 0.000004) as ret,
-        any(tr.fees) as fees,
-        argMax(tm.total_volume_usd, tm.updated_at) as vol,
-        argMax(tm.unique_holders, tm.updated_at) as holders
-    FROM token_metrics tm
-    LEFT JOIN (
-        SELECT 
-            base_address, 
-            sum(priority_fee + coin_creator_fee) as fees
-        FROM trades
-        GROUP BY base_address
-    ) tr ON tm.token_address = tr.base_address
-    GROUP BY tm.token_address
-    HAVING ret <= 10000
+    Internal logic: returns (buckets_dict, thresholds_dict, count_manipulated)
+    buckets_dict: {class_id: [list of tokens]}
     """
-    rows = client.execute(query)
-    
     # 1. Initial Classification
-    temp_map = {} # token -> {class_id, fees, vol, holders}
-    
-    # Storage for stats calculation
-    class_stats = {i: {'fees': [], 'vol': [], 'holders': []} for i in range(len(RETURN_THRESHOLDS)-1)}
+    temp_buckets = {i: [] for i in range(len(RETURN_THRESHOLDS))} 
     
-    print(f"   -> Initial classification of {len(rows)} tokens...")
-    for r in rows:
-        token_addr = r[0]
-        ret_val = r[1]
-        fees = r[2] or 0.0
-        vol = r[3] or 0.0
-        holders = r[4] or 0
+    for d in data:
+        ret = d["ret"]
+        if ret > 10000: continue 
         
-        class_id = -1
+        cid = 0 
+        found = False
         for i in range(len(RETURN_THRESHOLDS) - 1):
             lower = RETURN_THRESHOLDS[i]
             upper = RETURN_THRESHOLDS[i+1]
-            if ret_val >= lower and ret_val < upper:
-                class_id = i
+            if ret >= lower and ret < upper:
+                cid = i
+                found = True
                 break
         
-        if class_id != -1:
-            temp_map[token_addr] = {'id': class_id, 'fees': fees, 'vol': vol, 'holders': holders}
-            class_stats[class_id]['fees'].append(fees)
-            class_stats[class_id]['vol'].append(vol)
-            class_stats[class_id]['holders'].append(holders)
+        if found:
+            d["class_id_initial"] = cid
+            temp_buckets[cid].append(d)
+        else:
+            if ret >= 10000: continue
+            d["class_id_initial"] = 0
+            temp_buckets[0].append(d)
 
-    # 2. Calculate Medians & Thresholds
+    # 2. Calculate Thresholds (50% of Median)
+    print("\n   -> Calculating Class Medians & Thresholds (Dynamic Outlier Detection)...")
     thresholds = {}
-    print("   -> Calculating Class Medians & Thresholds (< 10% of Median)...")
-    for i in range(1, 5): # Check classes 1, 2, 3, 4 (Profitable to PVE)
-        # Class 0 (Garbage) is not checked/filtered
-        if len(class_stats[i]['fees']) > 0:
-            med_fees = np.median(class_stats[i]['fees'])
-            med_vol = np.median(class_stats[i]['vol'])
-            med_holders = np.median(class_stats[i]['holders'])
+    
+    for i in range(1, len(RETURN_THRESHOLDS)-1):
+        items = temp_buckets.get(i, [])
+        if len(items) > 5:
+            fees = [x["fees_sol"] for x in items]
+            vols = [x["volume_usd"] for x in items]
+            holders = [x["unique_holders"] for x in items]
+            
+            med_fees = np.median(fees)
+            med_vol = np.median(vols)
+            med_holders = np.median(holders)
             
             thresholds[i] = {
                 'fees': med_fees * 0.5,
                 'vol': med_vol * 0.5,
                 'holders': med_holders * 0.5
             }
-            print(f"      [Class {i}] Median Fees: {med_fees:.4f} (Thresh: {thresholds[i]['fees']:.4f}) | Median Vol: ${med_vol:.0f} (Thresh: ${thresholds[i]['vol']:.0f}) | Median Holders: {med_holders:.0f} (Thresh: {thresholds[i]['holders']:.0f})")
         else:
-            thresholds[i] = {'fees': 0, 'vol': 0, 'holders': 0}
+             thresholds[i] = {'fees': 0, 'vol': 0, 'holders': 0}
 
     # 3. Reclassification
-    print("   -> Detecting Manipulated Outliers...")
-    final_map = {}
-    manipulated_count = 0
+    final_buckets = {i: [] for i in range(len(RETURN_THRESHOLDS))}
+    final_buckets[MANIPULATED_CLASS_ID] = [] 
+    
+    count_manipulated = 0
     
-    for token, data in temp_map.items():
-        cid = data['id']
-        # Only check if it's a "successful" class (ID > 0)
-        if cid > 0 and cid in thresholds:
-            t = thresholds[cid]
-            # Condition: If ANY metric is suspiciously low
-            is_manipulated = (data['fees'] < t['fees']) or (data['vol'] < t['vol']) or (data['holders'] < t['holders'])
+    for cid, items in temp_buckets.items():
+        for d in items:
+            final_cid = cid
+            if cid > 0 and cid in thresholds:
+                t = thresholds[cid]
+                if (d["fees_sol"] < t['fees']) or (d["volume_usd"] < t['vol']) or (d["unique_holders"] < t['holders']):
+                    final_cid = MANIPULATED_CLASS_ID
+                    count_manipulated += 1
             
-            if is_manipulated:
-                final_map[token] = MANIPULATED_CLASS_ID
-                manipulated_count += 1
-            else:
-                final_map[token] = cid
-        else:
-            final_map[token] = cid
+            d["class_id_final"] = final_cid
+            if final_cid not in final_buckets:
+                final_buckets[final_cid] = []
+            final_buckets[final_cid].append(d)
             
-    print(f"   -> Reclassification Complete. identified {manipulated_count} manipulated tokens.")
-    return final_map, thresholds
-
-def analyze():
-    client = get_client()
+    return final_buckets, thresholds, count_manipulated
 
-    print("=== SEGMENTED DISTRIBUTION ANALYSIS ===")
-    
-    # 1. Get Classified Map AND Thresholds
-    class_map, thresholds = get_return_class_map(client)
+def get_return_class_map(client):
+    """
+    Returns (map {token_addr: class_id}, thresholds)
+    Used by cache_dataset.py
+    """
+    data = fetch_all_metrics(client)
+    buckets, thresholds, _ = _classify_tokens(data)
     
-    # 2. Invert Map for easy lookups (still useful for counts or smaller segments)
-    segments_tokens = {}
-    for t, c in class_map.items():
-        if c not in segments_tokens:
-            segments_tokens[c] = []
-        segments_tokens[c].append(t)
+    # Flatten buckets to map
+    ret_map = {}
+    for cid, items in buckets.items():
+        for d in items:
+            ret_map[d["token_address"]] = cid
+            
+    return ret_map, thresholds
+
+def print_stats(name, values):
+    """
+    prints compact stats: mean, p50, p90, p99
+    """
+    if not values:
+        print(f"  {name}: No data")
+        return
         
-    # Define Labels from thresholds so bucket changes don't silently desync output.
-    labels = {}
-    for i in range(len(RETURN_THRESHOLDS) - 1):
-        lower = RETURN_THRESHOLDS[i]
-        upper = RETURN_THRESHOLDS[i + 1]
-        labels[i] = f"{i}. {lower}x - {upper}x"
-    labels[MANIPULATED_CLASS_ID] = f"{MANIPULATED_CLASS_ID}. MANIPULATED (Fake Metrics)"
-    
-    # Common SQL parts
-    # We need a robust base for the WHERE clause variables (fees, vol, holders)
-    # Since we can't easily alias in the WHERE clause of a subquery filter without re-joining,
-    # we will rely on a standardized CTE-like structure or just simpler subqueries in the condition.
+    vals = np.array(values)
+    mean = np.mean(vals)
+    p50 = np.percentile(vals, 50)
+    p90 = np.percentile(vals, 90)
+    p99 = np.percentile(vals, 99)
+    nonzero = np.count_nonzero(vals)
+    nonzero_rate = nonzero / len(vals)
     
-    # Efficient Token Metrics View
-    # We need to filter based on: ret, fees, vol, holders
-    # fees come from trades (sum), vol/holders/ret from token_metrics (argMax)
+    print(f"  {name}: mean={mean:.4f} p50={p50:.4f} p90={p90:.4f} p99={p99:.4f} nonzero_rate={nonzero_rate:.3f} (n={len(vals)})")
+
+def analyze():
+    client = get_client()
+    data = fetch_all_metrics(client)
+    final_buckets, thresholds, count_manipulated = _classify_tokens(data)
     
-    # To keep query size small, we define the criteria logic in SQL.
-    # But we need 'fees' which is an aggregate.
-    # So we define a base cohort query that computes these 4 values for EVERY token, 
-    # and then wrap it with the WHERE clause.
+    print(f"   -> Reclassification Complete. Identified {count_manipulated} manipulated tokens.")
+    print("\n=== SEGMENTED DISTRIBUTION ANALYSIS ===")
     
-    base_cohort_source = """
-        SELECT 
-            tm.token_address as join_key,
-            (argMax(tm.ath_price_usd, tm.updated_at) / 0.000004) as ret,
-            any(tr.fees) as fees,
-            argMax(tm.total_volume_usd, tm.updated_at) as vol,
-            argMax(tm.unique_holders, tm.updated_at) as holders
-        FROM token_metrics tm
-        LEFT JOIN (
-            SELECT base_address, sum(priority_fee + coin_creator_fee) as fees
-            FROM trades
-            GROUP BY base_address
-        ) tr ON tm.token_address = tr.base_address
-        GROUP BY tm.token_address
-    """
+    # Print Thresholds debug
+    for k, t in thresholds.items():
+         if t['fees'] > 0:
+             print(f"      [Class {k}] Thresh: Fees>{t['fees']:.3f} Vol>${t['vol']:.0f} Holders>{t['holders']:.0f}")
 
-    # Iterate through known classes
-    for cid in sorted(labels.keys()):
-        label = labels[cid]
-        tokens = segments_tokens.get(cid, [])
-        count = len(tokens)
-        
-        print(f"\n\n==================================================")
-        print(f"SEGMENT: {label}")
-        print(f"==================================================")
-        print(f"Tokens in segment: {count}")
-        
-        if count == 0:
-            continue
-            
-        # Construct SQL Condition based on ID
-        condition = "1=0" # Default fail
+    sorted_classes = sorted([k for k in final_buckets.keys() if k != MANIPULATED_CLASS_ID]) + [MANIPULATED_CLASS_ID]
+    
+    for cid in sorted_classes:
+        items = final_buckets.get(cid, [])
+        if not items: continue 
         
-        if cid == 0:
-            # Garbage: Just Return < 3. 
-            # Note: Technically it also includes tokens that might have been >3x but <10000x... 
-            # BUT our Python/Map logic says Garbage is class 0. 
-            # The only way to be class 0 in the map is if ret < 3.
-            # Downgraded tokens go to Class 5.
-            condition = "ret < 3"
-            
-        elif cid == MANIPULATED_CLASS_ID:
-            # Manipulated: 
-            # It's the collection of (Class K logic AND is_outlier)
-            sub_conds = []
-            for k in range(1, 5):
-                if k in thresholds:
-                    t = thresholds[k]
-                    # Range for Class K
-                    lower = RETURN_THRESHOLDS[k]
-                    upper = RETURN_THRESHOLDS[k+1]
-                    # Outlier logic
-                    sub_conds.append(f"(ret >= {lower} AND ret < {upper} AND (fees < {t['fees']} OR vol < {t['vol']} OR holders < {t['holders']}))")
-            
-            if sub_conds:
-                condition = " OR ".join(sub_conds)
-            
+        if cid == MANIPULATED_CLASS_ID:
+             label = f"{cid}. MANIPULATED / FAKE (Outliers from {1}~{4})"
+        elif cid < len(RETURN_THRESHOLDS)-1:
+             label = f"{cid}. {RETURN_THRESHOLDS[cid]}x - {RETURN_THRESHOLDS[cid+1]}x"
         else:
-            # Normal Classes 1-4
-            if cid in thresholds:
-                t = thresholds[cid]
-                lower = RETURN_THRESHOLDS[cid]
-                upper = RETURN_THRESHOLDS[cid+1]
-                # Valid logic: In Range AND NOT Outlier
-                condition = f"(ret >= {lower} AND ret < {upper} AND fees >= {t['fees']} AND vol >= {t['vol']} AND holders >= {t['holders']})"
+             label = f"{cid}. Unknown"
+             
+        print(f"\nSEGMENT: {label}")
+        print("="*50)
+        print(f"Tokens in segment: {len(items)}")
         
-        # Final Cohort SQL: Select keys satisfying the condition
-        # We wrap the base source
-        cohort_sql = f"""
-            SELECT join_key FROM (
-                {base_cohort_source}
-            ) WHERE {condition}
-        """
-
-        # Helper to construct the full condition "join_key IN (...)"
-        # NOW we use the subquery instead of a literal list
-        def make_query(inner, cohort_subquery):
-             return f"""
-                SELECT * FROM (
-                    {inner}
-                ) WHERE join_key IN (
-                    {cohort_subquery}
-                )
-             """
-
-        # --- Metrics Definitions ---
+        bundled = [x["bundled_pct"] for x in items]
+        dev_hold = [x["dev_hold_pct"] for x in items]
+        fees = [x["fees_sol"] for x in items]
+        snipers = [x["snipers_pct"] for x in items]
         
-        # 1. Fees (SOL)
-        fees_inner = """
-            SELECT base_address as join_key, sum(priority_fee + coin_creator_fee) as val
-            FROM trades
-            GROUP BY base_address
-        """
-        fees_buckets = """
-            case
-                when val < 0.001 then '1. < 0.001 SOL'
-                when val >= 0.001 AND val < 0.01 then '2. 0.001 - 0.01'
-                when val >= 0.01 AND val < 0.1 then '3. 0.01 - 0.1'
-                when val >= 0.1 AND val < 1 then '4. 0.1 - 1'
-                when val >= 1 then '5. > 1 SOL'
-                else 'Unknown'
-            end
-        """
-        print_distribution_stats(client, "Total Fees (SOL)", make_query(fees_inner, cohort_sql), fees_buckets)
-
-        # 2. Volume (USD)
-        vol_inner = """
-            SELECT base_address as join_key, sum(total_usd) as val
-            FROM trades
-            GROUP BY base_address
-        """
-        vol_buckets = """
-            case
-                when val < 1000 then '1. < $1k'
-                when val >= 1000 AND val < 10000 then '2. $1k - $10k'
-                when val >= 10000 AND val < 100000 then '3. $10k - $100k'
-                when val >= 100000 AND val < 1000000 then '4. $100k - $1M'
-                when val >= 1000000 then '5. > $1M'
-                else 'Unknown'
-            end
-        """
-        print_distribution_stats(client, "Total Volume (USD)", make_query(vol_inner, cohort_sql), vol_buckets)
-
-        # 3. Unique Holders
-        holders_inner = """
-            SELECT token_address as join_key, argMax(unique_holders, updated_at) as val
-            FROM token_metrics
-            GROUP BY token_address
-        """
-        holders_buckets = """
-            case
-                when val < 10 then '1. < 10'
-                when val >= 10 AND val < 50 then '2. 10 - 50'
-                when val >= 50 AND val < 100 then '3. 50 - 100'
-                when val >= 100 AND val < 500 then '4. 100 - 500'
-                when val >= 500 then '5. > 500'
-                else 'Unknown'
-            end
-        """
-        print_distribution_stats(client, "Unique Holders", make_query(holders_inner, cohort_sql), holders_buckets)
-
-        # 4. Snipers % Supply
-        snipers_inner = """
-            SELECT
-                m.base_address as join_key,
-                (m.val / t.total_supply * 100) as val
-            FROM (
-                SELECT
-                    base_address,
-                    sumIf(base_amount, buyer_rank <= 70) as val
-                FROM (
-                    SELECT
-                        base_address,
-                        base_amount,
-                        dense_rank() OVER (PARTITION BY base_address ORDER BY min_slot, min_idx) as buyer_rank
-                    FROM (
-                        SELECT 
-                            base_address, 
-                            maker, 
-                            min(slot) as min_slot, 
-                            min(transaction_index) as min_idx,
-                            sum(base_amount) as base_amount 
-                        FROM trades
-                        WHERE trade_type = 0
-                        GROUP BY base_address, maker
-                    )
-                )
-                GROUP BY base_address
-            ) m
-            JOIN (
-                SELECT token_address, argMax(total_supply, updated_at) as total_supply
-                FROM tokens
-                GROUP BY token_address
-            ) t ON m.base_address = t.token_address
-            WHERE t.total_supply > 0
-        """
-        pct_buckets = """
-            case
-                when val < 1 then '1. < 1%'
-                when val >= 1 AND val < 5 then '2. 1% - 5%'
-                when val >= 5 AND val < 10 then '3. 5% - 10%'
-                when val >= 10 AND val < 20 then '4. 10% - 20%'
-                when val >= 20 AND val < 50 then '5. 20% - 50%'
-                when val >= 50 then '6. > 50%'
-                else 'Unknown'
-            end
-        """
-        print_distribution_stats(client, "Snipers % Supply (Top 70)", make_query(snipers_inner, cohort_sql), pct_buckets)
-
-        # 5. Bundled % Supply
-        bundled_inner = """
-            SELECT
-                m.base_address as join_key,
-                (m.val / t.total_supply * 100) as val
-            FROM (
-                SELECT
-                    t.base_address,
-                    sum(t.base_amount) as val
-                FROM trades t
-                JOIN (
-                    SELECT base_address, min(slot) as min_slot
-                    FROM trades
-                    GROUP BY base_address
-                ) m ON t.base_address = m.base_address AND t.slot = m.min_slot
-                WHERE t.trade_type = 0
-                GROUP BY t.base_address
-            ) m
-            JOIN (
-                SELECT token_address, argMax(total_supply, updated_at) as total_supply
-                FROM tokens
-                GROUP BY token_address
-            ) t ON m.base_address = t.token_address
-            WHERE t.total_supply > 0
-        """
-        print_distribution_stats(client, "Bundled % Supply", make_query(bundled_inner, cohort_sql), pct_buckets)
-
-        # 6. Dev Holding % Supply
-        dev_inner = """
-            SELECT 
-                t.token_address as join_key,
-                (wh.current_balance / (t.total_supply / pow(10, t.decimals)) * 100) as val
-            FROM (
-                 SELECT token_address, argMax(creator_address, updated_at) as creator_address, argMax(total_supply, updated_at) as total_supply, argMax(decimals, updated_at) as decimals
-                 FROM tokens 
-                 GROUP BY token_address
-            ) t
-            JOIN (
-               SELECT mint_address, wallet_address, argMax(current_balance, updated_at) as current_balance
-               FROM wallet_holdings
-               GROUP BY mint_address, wallet_address
-            ) wh ON t.token_address = wh.mint_address AND t.creator_address = wh.wallet_address
-            WHERE t.total_supply > 0
-        """
-        print_distribution_stats(client, "Dev Holding % Supply", make_query(dev_inner, cohort_sql), pct_buckets)
-
-
-
-        # 8. Time to ATH (Seconds)
-        time_ath_inner = """
-            SELECT
-                base_address as join_key,
-                (argMax(timestamp, price_usd) - min(timestamp)) as val
-            FROM trades
-            GROUP BY base_address
-        """
-        time_ath_buckets = """
-            case
-                when val < 5 then '1. < 5s'
-                when val >= 5 AND val < 30 then '2. 5s - 30s'
-                when val >= 30 AND val < 60 then '3. 30s - 1m'
-                when val >= 60 AND val < 300 then '4. 1m - 5m'
-                when val >= 300 AND val < 3600 then '5. 5m - 1h'
-                when val >= 3600 then '6. > 1h'
-                else 'Unknown'
-            end
-        """
-        print_distribution_stats(client, "Time to ATH (Seconds)", make_query(time_ath_inner, cohort_sql), time_ath_buckets)
-
+        print_stats("bundled_pct", bundled)
+        print_stats("dev_hold_pct", dev_hold)
+        print_stats("fees_sol", fees)
+        print_stats("snipers_pct", snipers)
 
 if __name__ == "__main__":
     analyze()

scripts/cache_dataset.py

@@ -6,6 +6,7 @@ import numpy as np
 import datetime
 import torch
 import json
+import math
 from pathlib import Path
 from tqdm import tqdm
 from dotenv import load_dotenv
@@ -23,6 +24,8 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 from data.data_loader import OracleDataset
 from data.data_fetcher import DataFetcher
 from scripts.analyze_distribution import get_return_class_map
+# Import quality score calculator
+from scripts.compute_quality_score import get_token_quality_scores, fetch_token_metrics, _bucket_id, _midrank_percentiles, EPS
 
 from clickhouse_driver import Client as ClickHouseClient
 from neo4j import GraphDatabase
@@ -94,6 +97,86 @@ def compute_save_ohlc_stats(client: ClickHouseClient, output_path: str):
         print(f"ERROR: Failed to compute OHLC stats: {e}")
         # Don't crash, let it try to proceed (though dataset might complain if file missing)
 
+def build_quality_missing_reason_map(client: ClickHouseClient, max_ret: float = 1e9):
+    """
+    Build a map: token_address -> reason string for why a quality score is missing.
+    This mirrors compute_quality_scores filtering and feature availability.
+    """
+    data = fetch_token_metrics(client)
+    metrics_by_token = {d.get("token_address"): d for d in data if d.get("token_address")}
+
+    # Build buckets with the same return filtering as compute_quality_scores
+    buckets = {}
+    for d in data:
+        ret_val = d.get("ret")
+        if ret_val is None or ret_val <= 0 or ret_val > max_ret:
+            continue
+        b = _bucket_id(ret_val)
+        if b == -1:
+            continue
+        d["bucket_id"] = b
+        buckets.setdefault(b, []).append(d)
+
+    # Same feature definitions as compute_quality_scores
+    feature_defs = [
+        ("fees_log", lambda d: math.log1p(d["fees_sol"]) if d.get("fees_sol") is not None else None, True),
+        ("volume_log", lambda d: math.log1p(d["volume_usd"]) if d.get("volume_usd") is not None else None, True),
+        ("holders_log", lambda d: math.log1p(d["unique_holders"]) if d.get("unique_holders") is not None else None, True),
+        ("time_to_ath_log", lambda d: math.log1p(d["time_to_ath_sec"]) if d.get("time_to_ath_sec") is not None else None, True),
+        ("fees_per_volume", lambda d: (d["fees_sol"] / (d["volume_usd"] + EPS)) if d.get("fees_sol") is not None and d.get("volume_usd") is not None else None, True),
+        ("fees_per_trade", lambda d: (d["fees_sol"] / (d["n_trades"] + EPS)) if d.get("fees_sol") is not None and d.get("n_trades") is not None else None, True),
+        ("holders_per_trade", lambda d: (d["unique_holders"] / (d["n_trades"] + EPS)) if d.get("unique_holders") is not None and d.get("n_trades") is not None else None, True),
+        ("holders_per_volume", lambda d: (d["unique_holders"] / (d["volume_usd"] + EPS)) if d.get("unique_holders") is not None and d.get("volume_usd") is not None else None, True),
+        ("snipers_pct", lambda d: d.get("snipers_pct"), True),
+        ("bundled_pct", lambda d: d.get("bundled_pct"), True),
+        ("dev_hold_pct", lambda d: d.get("dev_hold_pct"), True),
+    ]
+
+    # Precompute percentiles per bucket + feature
+    bucket_feature_percentiles = {}
+    for b, items in buckets.items():
+        feature_percentiles = {}
+        for fname, fget, _pos in feature_defs:
+            vals = []
+            for d in items:
+                v = fget(d)
+                if v is None or (isinstance(v, float) and (math.isnan(v) or math.isinf(v))):
+                    continue
+                vals.append((d["token_address"], v))
+            feature_percentiles[fname] = _midrank_percentiles(vals)
+        bucket_feature_percentiles[b] = feature_percentiles
+
+    def _reason_for(token_address: str) -> str:
+        d = metrics_by_token.get(token_address)
+        if not d:
+            return "no metrics found (missing from token_metrics/trades/mints joins)"
+        ret_val = d.get("ret")
+        if ret_val is None:
+            return "ret is None (missing ATH/launch metrics)"
+        if ret_val <= 0:
+            return f"ret <= 0 ({ret_val})"
+        if ret_val > max_ret:
+            return f"ret > max_ret ({ret_val} > {max_ret})"
+        b = _bucket_id(ret_val)
+        if b == -1:
+            return f"ret {ret_val} not in RETURN_THRESHOLDS"
+        items = buckets.get(b, [])
+        if not items:
+            return f"bucket {b} empty after filtering"
+        feature_percentiles = bucket_feature_percentiles.get(b, {})
+        has_any = False
+        missing_features = []
+        for fname, _fget, _pos in feature_defs:
+            if feature_percentiles.get(fname, {}).get(token_address) is None:
+                missing_features.append(fname)
+            else:
+                has_any = True
+        if not has_any:
+            return "no valid feature percentiles for token (all features missing/invalid)"
+        return f"unexpected: has feature percentiles but no score; missing features={','.join(missing_features)}"
+
+    return _reason_for
+
 def main():
     load_dotenv()
     
@@ -140,10 +223,15 @@ def main():
         data_fetcher = DataFetcher(clickhouse_client=clickhouse_client, neo4j_driver=neo4j_driver)
         
         # Pre-fetch the Return Class Map
-        # tokens not in this map (e.g. >10k x) are INVALID and will be skipped
         print("INFO: Fetching Return Classification Map...")
         return_class_map, thresholds = get_return_class_map(clickhouse_client)
         print(f"INFO: Loaded {len(return_class_map)} valid classified tokens.")
+
+        # Pre-fetch Quality Scores
+        print("INFO: Fetching Token Quality Scores...")
+        quality_scores_map = get_token_quality_scores(clickhouse_client)
+        quality_missing_reason = build_quality_missing_reason_map(clickhouse_client, max_ret=1e9)
+        print(f"INFO: Loaded {len(quality_scores_map)} quality scores.")
     
         dataset = OracleDataset(
             data_fetcher=data_fetcher,
@@ -158,67 +246,103 @@ def main():
         if len(dataset) == 0:
             print("WARNING: Dataset initialization resulted in 0 samples. Nothing to cache.")
             return
+
+        # --- FILTER DATASET BY CLASS MAP ---
+        # Only keep mints that are classified (valid return, sufficient data)
+        original_size = len(dataset)
+        print(f"INFO: Filtering dataset... Original size: {original_size}")
+        dataset.sampled_mints = [
+            m for m in dataset.sampled_mints 
+            if m['mint_address'] in return_class_map
+        ]
+        filtered_size = len(dataset)
+        filtered_count = original_size - filtered_size
+        print(f"INFO: Filtered size: {filtered_size}")
+
+        if len(dataset) == 0:
+            print("WARNING: No tokens remain after filtering by return_class_map.")
+            return
     
         # --- 3. Iterate and cache each item ---
         print(f"INFO: Starting to generate and cache {len(dataset)} samples...")
         
-        metadata_path = output_dir / "metadata.jsonl"
-        print(f"INFO: Writing metadata to {metadata_path}")
-        
         skipped_count = 0
-        filtered_count = 0
         cached_count = 0
         
-        # Open metadata file in append mode 
-        with open(metadata_path, 'a') as meta_f:
-            for i in tqdm(range(len(dataset)), desc="Caching samples"):
-                mint_addr = dataset.sampled_mints[i]['mint_address']
-                
-                # 1. Filter Check
-                if mint_addr not in return_class_map:
-                    # Token is effectively "filtered out" (e.g. > 10,000x return or missing metrics)
-                    filtered_count += 1
-                    continue
-                
-                class_id = return_class_map[mint_addr]
-    
-                try:
-                    item = dataset.__cacheitem__(i)
-                    if item is None:
-                        skipped_count += 1
-                        continue
-                    
-                    filename = f"sample_{i}.pt"
-                    output_path = output_dir / filename
-                    torch.save(item, output_path)
-                    
-                    # Write metadata entry
-                    # Minimizing IO overhead by keeping line short
-                    meta_entry = {"file": filename, "class_id": class_id}
-                    meta_f.write(json.dumps(meta_entry) + "\n")
-                    
-                    cached_count += 1
-                    
-                except Exception as e:
-                    error_msg = str(e)
-                    # If a FATAL error occurs (e.g. persistent DB auth failure), stop the script immediately.
-                    if "FATAL" in error_msg or "AuthenticationRateLimit" in error_msg:
-                        print(f"\nCRITICAL: Fatal error encountered processing sample {i}. Stopping execution.\nError: {e}", file=sys.stderr)
-                        sys.exit(1)
-                    
-                    print(f"\nERROR: Failed to generate or save sample {i} for mint '{mint_addr}'. Error: {e}", file=sys.stderr)
-                    # print trackback
-                    import traceback
-                    traceback.print_exc()
+        for i in tqdm(range(len(dataset)), desc="Caching samples"):
+            mint_addr = dataset.sampled_mints[i]['mint_address']
+            
+            # (No need to check if in return_class_map anymore, we filtered)
+            class_id = return_class_map[mint_addr]
+
+            try:
+                item = dataset.__cacheitem__(i)
+                if item is None:
                     skipped_count += 1
                     continue
+                
+                # Require quality score only for samples that will be cached
+                if mint_addr not in quality_scores_map:
+                    reason = quality_missing_reason(mint_addr)
+                    raise RuntimeError(
+                        f"Missing quality score for mint {mint_addr}. Reason: {reason}. "
+                        "Refusing to cache without quality_score."
+                    )
+                q_score = quality_scores_map[mint_addr]
+
+                # INJECT QUALITY SCORE INTO TENSOR DICT
+                item["quality_score"] = q_score
+                item["class_id"] = class_id
+                
+                filename = f"sample_{i}.pt"
+                output_path = output_dir / filename
+                torch.save(item, output_path)
+                
+                cached_count += 1
+                
+                # Log progress details (reflect all cached event lists)
+                n_trades = len(item.get("trades", []))
+                n_transfers = len(item.get("transfers", []))
+                n_pool_creations = len(item.get("pool_creations", []))
+                n_liquidity_changes = len(item.get("liquidity_changes", []))
+                n_fee_collections = len(item.get("fee_collections", []))
+                n_burns = len(item.get("burns", []))
+                n_supply_locks = len(item.get("supply_locks", []))
+                n_migrations = len(item.get("migrations", []))
+                n_ohlc = len(item.get("ohlc_1s", [])) if item.get("ohlc_1s") is not None else 0
+                n_snapshots_5m = len(item.get("snapshots_5m", []))
+                n_holders = len(item.get("holder_snapshots_list", []))
+                
+                tqdm.write(f"  + Cached: {mint_addr} | Class: {class_id} | Q: {q_score:.4f}")
+                tqdm.write(
+                    "    Events | "
+                    f"Trades: {n_trades} | Transfers: {n_transfers} | Pool Creations: {n_pool_creations} | "
+                    f"Liquidity Changes: {n_liquidity_changes} | Fee Collections: {n_fee_collections} | "
+                    f"Burns: {n_burns} | Supply Locks: {n_supply_locks} | Migrations: {n_migrations}"
+                )
+                tqdm.write(
+                    f"    Derived | Mint: 1 | Ohlc 1s: {n_ohlc} | Snapshots 5m: {n_snapshots_5m} | Holder Snapshots: {n_holders}"
+                )
+                
+            except Exception as e:
+                error_msg = str(e)
+                # If a FATAL error occurs (e.g. persistent DB auth failure), stop the script immediately.
+                if "FATAL" in error_msg or "AuthenticationRateLimit" in error_msg:
+                    print(f"\nCRITICAL: Fatal error encountered processing sample {i}. Stopping execution.\nError: {e}", file=sys.stderr)
+                    sys.exit(1)
+                
+                print(f"\nERROR: Failed to generate or save sample {i} for mint '{mint_addr}'. Error: {e}", file=sys.stderr)
+                # print trackback
+                import traceback
+                traceback.print_exc()
+                skipped_count += 1
+                continue
         
         print(f"\n--- Caching Complete ---")
         print(f"Successfully cached: {cached_count} items.")
         print(f"Filtered (Invalid/High Return): {filtered_count} items.")
         print(f"Skipped (Errors/Empty): {skipped_count} items.")
         print(f"Cache location: {output_dir.resolve()}")
-        print(f"Metadata location: {metadata_path.resolve()}")
 
     finally:
         # --- 4. Close connections ---

scripts/compute_quality_score.py

@@ -87,15 +87,15 @@ def fetch_token_metrics(client) -> List[dict]:
             FROM trades
             GROUP BY base_address
         ),
-        -- 2. Token metadata (supply, decimals, creator)
+        -- 2. "Token list derived MINTS.
         token_meta_raw AS (
             SELECT
-                token_address,
-                argMax(creator_address, updated_at) AS creator_address,
-                argMax(total_supply, updated_at) AS total_supply,
-                argMax(decimals, updated_at) AS decimals
-            FROM tokens
-            GROUP BY token_address
+                mint_address AS token_address,
+                argMax(creator_address, timestamp) AS creator_address,
+                argMax(total_supply, timestamp) AS total_supply,
+                argMax(token_decimals, timestamp) AS decimals
+            FROM mints
+            GROUP BY mint_address
         ),
         token_meta AS (
             SELECT
@@ -161,28 +161,21 @@ def fetch_token_metrics(client) -> List[dict]:
             GROUP BY s.base_address
         ),
 
-        -- 6. BUNDLED: Identify bundled addresses, sum their PEAK balances
-        --    Bundled definition: Bought in the same slot as the very first buy slot for that token.
-        bundled_list AS (
-            SELECT
-                t.base_address,
-                t.maker
-            FROM trades t
-            JOIN (
-                SELECT base_address, min(slot) AS min_slot
-                FROM trades
-                GROUP BY base_address
-            ) m ON t.base_address = m.base_address AND t.slot = m.min_slot
-            WHERE t.trade_type = 0 -- buy
-            GROUP BY t.base_address, t.maker
-        ),
+        -- 6. BUNDLED: Sum the base_amount of ALL trades that happened in a slot with multiple buys
         bundled_agg AS (
             SELECT
-                b.base_address AS token_address,
-                sum(wp.peak_balance) AS bundled_total_peak
-            FROM bundled_list b
-            JOIN wallet_peaks wp ON b.base_address = wp.mint_address AND b.maker = wp.wallet_address
-            GROUP BY b.base_address
+                t.base_address AS token_address,
+                sum(t.base_amount) AS bundled_total_peak
+            FROM trades t
+            WHERE (t.base_address, t.slot) IN (
+                 SELECT base_address, slot
+                 FROM trades
+                 WHERE trade_type = 0 -- buy
+                 GROUP BY base_address, slot
+                 HAVING count() > 1
+            )
+            AND t.trade_type = 0 -- buy
+            GROUP BY t.base_address
         ),
 
         -- 7. DEV HOLD: Creator's Peak Balance
@@ -196,7 +189,7 @@ def fetch_token_metrics(client) -> List[dict]:
         )
 
     SELECT
-        r.token_address,
+        t.token_address,
         r.ret,
         r.unique_holders,
         f.fees_sol,
@@ -205,14 +198,14 @@ def fetch_token_metrics(client) -> List[dict]:
         (f.t_ath - f.t0) AS time_to_ath_sec,
         -- Calculate Percentages using Peak Sums / Total Supply
         (COALESCE(s.snipers_total_peak, 0) / t.adj_supply * 100) AS snipers_pct,
-        (COALESCE(b.bundled_total_peak, 0) / t.adj_supply * 100) AS bundled_pct,
+        (COALESCE(b.bundled_total_peak, 0) / t.total_supply * 100) AS bundled_pct,
         (COALESCE(d.dev_peak, 0)           / t.adj_supply * 100) AS dev_hold_pct
-    FROM ret_agg r
-    JOIN token_meta t ON r.token_address = t.token_address
-    LEFT JOIN trade_agg f ON r.token_address = f.base_address
-    LEFT JOIN snipers_agg s ON r.token_address = s.token_address
-    LEFT JOIN bundled_agg b ON r.token_address = b.token_address
-    LEFT JOIN dev_hold_agg d ON r.token_address = d.token_address
+    FROM token_meta t
+    LEFT JOIN ret_agg r ON t.token_address = r.token_address
+    LEFT JOIN trade_agg f ON t.token_address = f.base_address
+    LEFT JOIN snipers_agg s ON t.token_address = s.token_address
+    LEFT JOIN bundled_agg b ON t.token_address = b.token_address
+    LEFT JOIN dev_hold_agg d ON t.token_address = d.token_address
     """
     rows = client.execute(query)
     cols = [
@@ -233,7 +226,7 @@ def fetch_token_metrics(client) -> List[dict]:
     return out
 
 
-def _compute_quality_scores(
+def compute_quality_scores(
     client,
     max_ret: float = 10000.0,
     rerank: bool = True,
@@ -251,12 +244,12 @@ def _compute_quality_scores(
         ("fees_per_trade", lambda d: (d["fees_sol"] / (d["n_trades"] + EPS)) if d["fees_sol"] is not None and d["n_trades"] is not None else None, True),
         ("holders_per_trade", lambda d: (d["unique_holders"] / (d["n_trades"] + EPS)) if d["unique_holders"] is not None and d["n_trades"] is not None else None, True),
         ("holders_per_volume", lambda d: (d["unique_holders"] / (d["volume_usd"] + EPS)) if d["unique_holders"] is not None and d["volume_usd"] is not None else None, True),
-        ("snipers_pct", lambda d: d["snipers_pct"], False),
-        ("bundled_pct", lambda d: d["bundled_pct"], False),
-        ("dev_hold_pct", lambda d: d["dev_hold_pct"], False),
+        ("snipers_pct", lambda d: d["snipers_pct"], True),
+        ("bundled_pct", lambda d: d["bundled_pct"], True),
+        ("dev_hold_pct", lambda d: d["dev_hold_pct"], True),
     ]
 
-    raw_metrics = ["snipers_pct", "bundled_pct", "dev_hold_pct"]
+    raw_metrics = ["snipers_pct", "bundled_pct", "dev_hold_pct", "fees_sol"] # Added fees_sol for diagnostic logging
 
     debug = None
     if with_debug:
@@ -357,6 +350,10 @@ def _compute_quality_scores(
                         "ret": d["ret"],
                         "q_raw": q_raw_map[t],
                         "q": q_final,
+                        # Pass through raw metrics for analysis
+                        "bundled_pct": d.get("bundled_pct"),
+                        "snipers_pct": d.get("snipers_pct"),
+                        "fees_sol": d.get("fees_sol"),
                     }
                 )
         else:
@@ -371,6 +368,10 @@ def _compute_quality_scores(
                         "ret": d["ret"],
                         "q_raw": q_raw_map[t],
                         "q": q_raw_map[t],
+                        # Pass through raw metrics for analysis
+                        "bundled_pct": d.get("bundled_pct"),
+                        "snipers_pct": d.get("snipers_pct"),
+                        "fees_sol": d.get("fees_sol"),
                     }
                 )
 
@@ -379,12 +380,7 @@ def _compute_quality_scores(
     return token_scores
 
 
-def compute_quality_scores(
-    client,
-    max_ret: float = 10000.0,
-    rerank: bool = True,
-) -> List[dict]:
-    return _compute_quality_scores(client, max_ret=max_ret, rerank=rerank, with_debug=False)
+
 
 
 def write_jsonl(path: str, rows: List[dict]) -> None:
@@ -491,6 +487,23 @@ def print_summary(scores: List[dict]) -> None:
             print(f"    Mean: {stats_q_raw['mean']:.4f} | Min: {stats_q_raw['min']:.4f} | Max: {stats_q_raw['max']:.4f}")
             print(f"    Q: p10={stats_q_raw['p10']:.2f} p50={stats_q_raw['p50']:.2f} p90={stats_q_raw['p90']:.2f} p99={stats_q_raw['p99']:.2f}")
 
+        # --- NEW: Print 3 Examples (Min, Mid, Max) ---
+        if items:
+            # Sort items by 'q' to find min/mid/max easily
+            items_sorted = sorted(items, key=lambda x: x.get("q", 0))
+            
+            ex_min = items_sorted[0]
+            ex_max = items_sorted[-1]
+            
+            # Find mid (closest to 0.0, or just median index? Request said "mean quality" which is 0.0)
+            # finding item with q closest to 0.0
+            ex_mid = min(items_sorted, key=lambda x: abs(x.get("q", 0) - 0.0))
+            
+            print("  Examples:")
+            print(f"    Low  (-1.0): {ex_min['token_address']} (q={ex_min.get('q',0):.4f}, ret={ex_min.get('ret',0):.2f}x)")
+            print(f"    Mid  (~0.0): {ex_mid['token_address']} (q={ex_mid.get('q',0):.4f}, ret={ex_mid.get('ret',0):.2f}x)")
+            print(f"    High ( 1.0): {ex_max['token_address']} (q={ex_max.get('q',0):.4f}, ret={ex_max.get('ret',0):.2f}x)")
+
 
 def print_diagnostics(debug: dict) -> None:
     if not debug:
@@ -563,6 +576,77 @@ def print_diagnostics(debug: dict) -> None:
             corr = _pearson_corr(xs, ys)
             print(f"  log(ret) vs {metric}: {corr:.4f} (n={len(xs)})")
 
+    # Removed placeholder
+    pass
+
+
+def print_high_ret_analysis(scores: List[dict]) -> None:
+    print("\n=== MID-HIGH RETURN SPLIT ANALYSIS (10x - 20x) ===")
+    
+    # 1. Filter for Mid-High Return Cohort (10x - 20x)
+    cohort = [s for s in scores if s.get("ret") is not None and s["ret"] >= 10.0 and s["ret"] < 20.0]
+    if not cohort:
+        print("No tokens 10x-20x found.")
+        return
+
+    print(f"Total tokens 10x-20x: {len(cohort)}")
+
+    # 2. Extract Bundled Pct
+    bundled_vals = [s.get("bundled_pct", 0) for s in cohort if s.get("bundled_pct") is not None]
+    if not bundled_vals:
+        print("No bundled_pct data found.")
+        return
+    
+    median_bundled = _percentile(sorted(bundled_vals), 0.50)
+    print(f"Median Bundled% for Cohort: {median_bundled:.2f}%")
+
+    # 3. Split
+    low_group = [s for s in cohort if (s.get("bundled_pct") or 0) <= median_bundled]
+    high_group = [s for s in cohort if (s.get("bundled_pct") or 0) > median_bundled]
+
+    # 4. Analyze Fees
+    def get_mean_fees(group):
+        fees = [s.get("fees_sol", 0) for s in group if s.get("fees_sol") is not None]
+        if not fees: return 0.0
+        return sum(fees) / len(fees)
+
+    mean_fees_low = get_mean_fees(low_group)
+    mean_fees_high = get_mean_fees(high_group)
+
+    print(f"\nGroup 1: LOW Bundled (<= {median_bundled:.2f}%)")
+    print(f"  Count: {len(low_group)}")
+    print(f"  Mean Fees: {mean_fees_low:.4f} SOL")
+
+    print(f"\nGroup 2: HIGH Bundled (> {median_bundled:.2f}%)")
+    print(f"  Count: {len(high_group)}")
+    print(f"  Mean Fees: {mean_fees_high:.4f} SOL")
+    
+    # Extra: Check returns too
+    def get_mean_ret(group):
+        rets = [s["ret"] for s in group]
+        if not rets: return 0.0
+        return sum(rets) / len(rets)
+        
+    print(f"  Mean Ret:  {get_mean_ret(high_group):.2f}x (vs Low: {get_mean_ret(low_group):.2f}x)")
+
+
+def get_token_quality_scores(client):
+    """
+    Returns a dictionary mapping token_address -> q (quality score)
+    """
+    # Force rerank=True to get final scores
+    results = compute_quality_scores(client, max_ret=1e9, rerank=True)
+    
+    # Return mapping
+    # If compute_quality_scores returns (scores, debug) tuple (when with_debug=True), handle it.
+    # Default call rerank=True returns 'scores' list if with_debug=False?
+    # No, looking at main, it returns 'scores' if no_diagnostics.
+    # But get_token_quality_scores uses default args. 
+    # Let's check compute_quality_score signature... it has with_debug=False default.
+    # So it returns 'scores'.
+    
+    return {r["token_address"]: r.get("q", 0.0) for r in results}
+
 
 def main():
     parser = argparse.ArgumentParser(description="Compute token quality/health score.")
@@ -577,7 +661,7 @@ def main():
         scores = compute_quality_scores(client, max_ret=args.max_ret, rerank=not args.no_rerank)
         debug = None
     else:
-        scores, debug = _compute_quality_scores(
+        scores, debug = compute_quality_scores(
             client,
             max_ret=args.max_ret,
             rerank=not args.no_rerank,
@@ -587,6 +671,7 @@ def main():
         print_summary(scores)
     if not args.no_diagnostics:
         print_diagnostics(debug)
+        print_high_ret_analysis(scores) # Call the new analysis
 
 
 if __name__ == "__main__":

token_stats.rs

@@ -0,0 +1,857 @@
+use crate::database::insert_rows;
+use crate::services::price_service::PriceService;
+use crate::types::{
+    EventPayload, EventType, MigrationRow, MintRow, TokenMetricsRow, TokenStaticRow, TradeRow,
+};
+use anyhow::{Context, Result, anyhow};
+use borsh::BorshDeserialize;
+use clickhouse::Client;
+use futures_util::future;
+use mpl_token_metadata::accounts::Metadata;
+use once_cell::sync::Lazy;
+use redis::aio::MultiplexedConnection;
+use redis::streams::{StreamReadOptions, StreamReadReply};
+use redis::{AsyncCommands, Client as RedisClient, FromRedisValue};
+use solana_client::nonblocking::rpc_client::RpcClient;
+use solana_program::program_pack::Pack;
+use solana_sdk::pubkey::Pubkey;
+use spl_token::state::Mint;
+use std::collections::{HashMap, HashSet};
+use std::env;
+use std::str::FromStr;
+use std::sync::Arc;
+use std::time::Duration;
+use tokio::sync::RwLock;
+
+type TokenCache = HashMap<String, TokenEntry>;
+
+fn env_parse<T: FromStr>(key: &str, default: T) -> T {
+    env::var(key)
+        .ok()
+        .and_then(|v| v.parse::<T>().ok())
+        .unwrap_or(default)
+}
+
+static TOKEN_STATS_CHUNK_SIZE: Lazy<usize> =
+    Lazy::new(|| env_parse("TOKEN_STATS_CHUNK_SIZE", 1000usize));
+
+#[derive(Debug, Clone)]
+struct TokenEntry {
+    token: TokenStaticRow,
+    metrics: TokenMetricsRow,
+}
+
+impl TokenEntry {
+    fn new(token: TokenStaticRow, metrics: Option<TokenMetricsRow>) -> Self {
+        let metrics = metrics
+            .unwrap_or_else(|| TokenMetricsRow::new(token.token_address.clone(), token.updated_at));
+        Self { token, metrics }
+    }
+}
+
+#[derive(Clone, Debug)]
+struct TokenContext {
+    timestamp: u32,
+    protocol: Option<u8>,
+    pool_address: Option<String>,
+    decimals: Option<u8>,
+}
+
+impl TokenContext {
+    fn new(
+        timestamp: u32,
+        protocol: Option<u8>,
+        pool_address: Option<String>,
+        decimals: Option<u8>,
+    ) -> Self {
+        Self {
+            timestamp,
+            protocol,
+            pool_address,
+            decimals,
+        }
+    }
+}
+
+fn record_token_context(
+    contexts: &mut HashMap<String, TokenContext>,
+    token_address: &str,
+    timestamp: u32,
+    protocol: Option<u8>,
+    pool_address: Option<String>,
+    decimals: Option<u8>,
+) {
+    if token_address.is_empty() {
+        return;
+    }
+
+    let mut pool_for_insert = pool_address.clone();
+    let entry = contexts
+        .entry(token_address.to_string())
+        .or_insert_with(|| {
+            TokenContext::new(timestamp, protocol, pool_for_insert.take(), decimals)
+        });
+
+    if timestamp < entry.timestamp {
+        entry.timestamp = timestamp;
+    }
+
+    if entry.protocol.is_none() {
+        entry.protocol = protocol;
+    }
+
+    let should_update_pool = entry
+        .pool_address
+        .as_ref()
+        .map(|p| p.is_empty())
+        .unwrap_or(true);
+    if should_update_pool {
+        if let Some(pool) = pool_address {
+            if !pool.is_empty() {
+                entry.pool_address = Some(pool);
+            }
+        }
+    }
+
+    if let Some(dec) = decimals {
+        entry.decimals = Some(dec);
+    }
+}
+
+fn pool_addresses_from_context(context: &TokenContext) -> Vec<String> {
+    context
+        .pool_address
+        .as_ref()
+        .filter(|addr| !addr.is_empty())
+        .map(|addr| vec![addr.clone()])
+        .unwrap_or_default()
+}
+
+fn event_success(event: &EventType) -> bool {
+    match event {
+        EventType::Trade(row) => row.success,
+        EventType::Mint(row) => row.success,
+        EventType::Migration(row) => row.success,
+        EventType::FeeCollection(row) => row.success,
+        EventType::Liquidity(row) => row.success,
+        EventType::PoolCreation(row) => row.success,
+        EventType::Transfer(row) => row.success,
+        EventType::SupplyLock(row) => row.success,
+        EventType::SupplyLockAction(row) => row.success,
+        EventType::Burn(row) => row.success,
+    }
+}
+
+pub struct TokenAggregator {
+    db_client: Client,
+    redis_conn: MultiplexedConnection,
+    rpc_client: Arc<RpcClient>,
+    price_service: PriceService,
+    backfill_mode: bool,
+}
+
+impl TokenAggregator {
+    pub async fn new(
+        db_client: Client,
+        redis_client: RedisClient,
+        rpc_client: Arc<RpcClient>,
+        price_service: PriceService,
+    ) -> Result<Self> {
+        let redis_conn = redis_client.get_multiplexed_async_connection().await?;
+        println!("[TokenAggregator] ✔️ Connected to ClickHouse, Redis, and Solana RPC.");
+
+        let backfill_mode =
+            env::var("BACKFILL_MODE").unwrap_or_else(|_| "false".to_string()) == "true";
+        Ok(Self {
+            db_client,
+            redis_conn,
+            rpc_client,
+            price_service,
+            backfill_mode,
+        })
+    }
+
+    pub async fn run(&mut self) -> Result<()> {
+        let stream_key = "event_queue";
+        let group_name = "token_aggregators";
+        let consumer_name = format!("consumer-tokens-{}", uuid::Uuid::new_v4());
+
+        let mut publisher_conn = self.redis_conn.clone();
+        let next_queue = "wallet_agg_queue";
+
+        let result: redis::RedisResult<()> = self
+            .redis_conn
+            .xgroup_create_mkstream(stream_key, group_name, "0")
+            .await;
+        if let Err(e) = result {
+            if !e.to_string().contains("BUSYGROUP") {
+                return Err(anyhow!(
+                    "[TokenAggregator] Failed to create consumer group: {}",
+                    e
+                ));
+            }
+            println!(
+                "[TokenAggregator] Consumer group '{}' already exists. Resuming.",
+                group_name
+            );
+        } else {
+            println!(
+                "[TokenAggregator] Created new consumer group '{}'.",
+                group_name
+            );
+        }
+
+        loop {
+            let messages = match self
+                .collect_events(stream_key, group_name, &consumer_name)
+                .await
+            {
+                Ok(msgs) => msgs,
+                Err(e) => {
+                    eprintln!(
+                        "[TokenAggregator] 🔴 Error reading from Redis: {}. Retrying...",
+                        e
+                    );
+                    tokio::time::sleep(Duration::from_secs(5)).await;
+                    continue;
+                }
+            };
+            if messages.is_empty() {
+                continue;
+            }
+
+            println!(
+                "[TokenAggregator] ⚙️ Starting processing for a new batch of {} events...",
+                messages.len()
+            );
+            let message_ids: Vec<String> = messages.iter().map(|(id, _)| id.clone()).collect();
+            let payloads: Vec<EventPayload> =
+                messages.into_iter().map(|(_, payload)| payload).collect();
+
+            match self.process_batch(payloads.clone()).await {
+                // Clone payloads to use them after processing
+                Ok(_) => {
+                    if !message_ids.is_empty() {
+                        // Forward each payload to the next queue in the pipeline
+                        for payload in payloads {
+                            let payload_data = bincode::serialize(&payload)?;
+                            let _: () = publisher_conn
+                                .xadd(next_queue, "*", &[("payload", payload_data)])
+                                .await?;
+                        }
+                        println!(
+                            "[TokenAggregator] ✅ Finished batch, forwarded {} events to {}.",
+                            message_ids.len(),
+                            next_queue
+                        );
+
+                        // Acknowledge the message from the source queue ('event_queue')
+                        let _: () = self
+                            .redis_conn
+                            .xack(stream_key, group_name, &message_ids)
+                            .await?;
+                        let _: i64 = self
+                            .redis_conn
+                            .xdel::<_, _, i64>(stream_key, &message_ids)
+                            .await?;
+                    }
+                }
+                Err(e) => {
+                    eprintln!(
+                        "[TokenAggregator] ❌ Failed to process batch, will not forward or ACK. Error: {}",
+                        e
+                    );
+                }
+            }
+        }
+    }
+
+    async fn process_batch(&self, payloads: Vec<EventPayload>) -> Result<()> {
+        let mut token_contexts: HashMap<String, TokenContext> = HashMap::new();
+        for payload in &payloads {
+            if !event_success(&payload.event) {
+                continue;
+            }
+            let decimals_map = &payload.token_decimals;
+            match &payload.event {
+                EventType::Trade(t) => {
+                    let pool = (!t.pool_address.is_empty()).then(|| t.pool_address.clone());
+                    record_token_context(
+                        &mut token_contexts,
+                        &t.base_address,
+                        t.timestamp,
+                        Some(t.protocol),
+                        pool.clone(),
+                        decimals_map.get(&t.base_address).cloned(),
+                    );
+                    record_token_context(
+                        &mut token_contexts,
+                        &t.quote_address,
+                        t.timestamp,
+                        Some(t.protocol),
+                        pool,
+                        decimals_map.get(&t.quote_address).cloned(),
+                    );
+                }
+                EventType::Mint(m) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &m.mint_address,
+                        m.timestamp,
+                        Some(m.protocol),
+                        (!m.pool_address.is_empty()).then(|| m.pool_address.clone()),
+                        Some(m.token_decimals),
+                    );
+                }
+                EventType::Migration(m) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &m.mint_address,
+                        m.timestamp,
+                        Some(m.protocol),
+                        (!m.pool_address.is_empty()).then(|| m.pool_address.clone()),
+                        decimals_map.get(&m.mint_address).cloned(),
+                    );
+                }
+                EventType::FeeCollection(f) => {
+                    let vault = (!f.vault_address.is_empty()).then(|| f.vault_address.clone());
+                    record_token_context(
+                        &mut token_contexts,
+                        &f.token_0_mint_address,
+                        f.timestamp,
+                        Some(f.protocol),
+                        vault.clone(),
+                        decimals_map.get(&f.token_0_mint_address).cloned(),
+                    );
+                    if let Some(token_1) = &f.token_1_mint_address {
+                        record_token_context(
+                            &mut token_contexts,
+                            token_1,
+                            f.timestamp,
+                            Some(f.protocol),
+                            vault.clone(),
+                            decimals_map.get(token_1).cloned(),
+                        );
+                    }
+                }
+                EventType::PoolCreation(p) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &p.base_address,
+                        p.timestamp,
+                        Some(p.protocol),
+                        (!p.pool_address.is_empty()).then(|| p.pool_address.clone()),
+                        p.base_decimals
+                            .or_else(|| decimals_map.get(&p.base_address).cloned()),
+                    );
+                    record_token_context(
+                        &mut token_contexts,
+                        &p.quote_address,
+                        p.timestamp,
+                        Some(p.protocol),
+                        (!p.pool_address.is_empty()).then(|| p.pool_address.clone()),
+                        p.quote_decimals
+                            .or_else(|| decimals_map.get(&p.quote_address).cloned()),
+                    );
+                }
+                EventType::Transfer(t) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &t.mint_address,
+                        t.timestamp,
+                        None,
+                        None,
+                        decimals_map.get(&t.mint_address).cloned(),
+                    );
+                }
+                EventType::SupplyLock(lock) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &lock.mint_address,
+                        lock.timestamp,
+                        Some(lock.protocol),
+                        None,
+                        decimals_map.get(&lock.mint_address).cloned(),
+                    );
+                }
+                EventType::SupplyLockAction(action) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &action.mint_address,
+                        action.timestamp,
+                        Some(action.protocol),
+                        None,
+                        decimals_map.get(&action.mint_address).cloned(),
+                    );
+                }
+                EventType::Burn(burn) => {
+                    record_token_context(
+                        &mut token_contexts,
+                        &burn.mint_address,
+                        burn.timestamp,
+                        None,
+                        None,
+                        decimals_map.get(&burn.mint_address).cloned(),
+                    );
+                }
+                EventType::Liquidity(_) => {}
+                _ => {}
+            }
+        }
+
+        if token_contexts.is_empty() {
+            println!("[TokenAggregator] -> Batch contains no relevant token events. Skipping.");
+            return Ok(());
+        }
+        println!(
+            "[TokenAggregator] -> Batch contains {} unique tokens.",
+            token_contexts.len()
+        );
+
+        let mut tokens = self
+            .fetch_tokens_from_db(&token_contexts.keys().cloned().collect::<Vec<_>>())
+            .await?;
+
+        let missing_tokens: Vec<String> = token_contexts
+            .keys()
+            .filter(|address| !tokens.contains_key(*address))
+            .cloned()
+            .collect();
+
+        if !missing_tokens.is_empty() {
+            println!(
+                "[TokenAggregator] -> Found {} new tokens to fetch metadata for.",
+                missing_tokens.len()
+            );
+
+            if !self.backfill_mode {
+                let fetch_futures = missing_tokens
+                    .iter()
+                    .map(|key| async move { (key.clone(), self.fetch_token_metadata(key).await) });
+                let fetched_results = future::join_all(fetch_futures).await;
+
+                for (key, rpc_result) in fetched_results {
+                    let context = match token_contexts.get(&key) {
+                        Some(ctx) => ctx.clone(),
+                        None => continue,
+                    };
+                    let protocol = context.protocol.unwrap_or(0);
+                    let token_row = match rpc_result {
+                        Ok((metadata, mint_data)) => {
+                            println!(
+                                "[TokenAggregator] -> ✅ Successfully fetched metadata for new token {}.",
+                                key
+                            );
+
+                            let creator = metadata
+                                .creators
+                                .as_ref()
+                                .and_then(|creators| creators.first())
+                                .map(|c| c.address.to_string())
+                                .unwrap_or_default();
+
+                            TokenStaticRow::new(
+                                key.clone(),
+                                context.timestamp,
+                                metadata.name.trim_end_matches('\0').to_string(),
+                                metadata.symbol.trim_end_matches('\0').to_string(),
+                                metadata.uri.trim_end_matches('\0').to_string(),
+                                mint_data.decimals,
+                                creator,
+                                pool_addresses_from_context(&context),
+                                protocol,
+                                mint_data.supply,
+                                metadata.is_mutable,
+                                Some(metadata.update_authority.to_string()),
+                                Option::from(mint_data.mint_authority)
+                                    .map(|pk: Pubkey| pk.to_string()),
+                                Option::from(mint_data.freeze_authority)
+                                    .map(|pk: Pubkey| pk.to_string()),
+                            )
+                        }
+                        Err(e) => {
+                            eprintln!(
+                                "[TokenAggregator] -> ❌ RPC failed for {}: {}. Creating placeholder.",
+                                key, e
+                            );
+                            TokenStaticRow::new(
+                                key.clone(),
+                                context.timestamp,
+                                String::new(),
+                                String::new(),
+                                String::new(),
+                                context.decimals.unwrap_or(0),
+                                String::new(),
+                                pool_addresses_from_context(&context),
+                                protocol,
+                                0,
+                                true,
+                                None,
+                                None,
+                                None,
+                            )
+                        }
+                    };
+                    tokens.insert(key.clone(), TokenEntry::new(token_row, None));
+                }
+            } else {
+                println!(
+                    "[TokenAggregator] -> Creating {} placeholder tokens in backfill mode.",
+                    missing_tokens.len()
+                );
+                for key in missing_tokens {
+                    if let Some(context) = token_contexts.get(&key) {
+                        let placeholder_row = TokenStaticRow::new(
+                            key.clone(),
+                            context.timestamp,
+                            String::new(),
+                            String::new(),
+                            String::new(),
+                            context.decimals.unwrap_or(0),
+                            String::new(),
+                            pool_addresses_from_context(context),
+                            context.protocol.unwrap_or(0),
+                            0,
+                            false,
+                            None,
+                            None,
+                            None,
+                        );
+                        tokens.insert(key.clone(), TokenEntry::new(placeholder_row, None));
+                    }
+                }
+            }
+        }
+
+        let trader_pairs_in_batch: Vec<(String, String)> = payloads
+            .iter()
+            .filter_map(|p| {
+                if let EventType::Trade(t) = &p.event {
+                    Some((t.base_address.clone(), t.maker.clone()))
+                } else {
+                    None
+                }
+            })
+            .collect();
+
+        let mut existing_traders = HashSet::new();
+        if !trader_pairs_in_batch.is_empty() {
+            for chunk in trader_pairs_in_batch.chunks(*TOKEN_STATS_CHUNK_SIZE) {
+                let mut cursor = self.db_client
+                    .query("SELECT DISTINCT (mint_address, wallet_address) FROM wallet_holdings WHERE (mint_address, wallet_address) IN ?")
+                    .bind(chunk)
+                    .fetch::<(String, String)>()?;
+
+                while let Some(pair) = cursor.next().await? {
+                    existing_traders.insert(pair);
+                }
+            }
+        }
+
+        let mut counted_in_this_batch: HashSet<(String, String)> = HashSet::new();
+
+        for payload in payloads.iter() {
+            if !event_success(&payload.event) {
+                continue;
+            }
+            match &payload.event {
+                EventType::Mint(mint) => self.process_mint(mint, &mut tokens),
+                EventType::Trade(trade) => {
+                    self.process_trade(
+                        trade,
+                        &mut tokens,
+                        &existing_traders,
+                        &mut counted_in_this_batch,
+                    );
+                }
+                EventType::Migration(migration) => self.process_migration(migration, &mut tokens),
+                _ => {}
+            }
+        }
+
+        self.finalize_and_persist(tokens).await
+    }
+
+    fn process_trade(
+        &self,
+        trade: &TradeRow,
+        tokens: &mut TokenCache,
+        existing_traders: &HashSet<(String, String)>,
+        counted_in_this_batch: &mut HashSet<(String, String)>,
+    ) {
+        if let Some(entry) = tokens.get_mut(&trade.base_address) {
+            entry.token.updated_at = trade.timestamp;
+            entry.metrics.updated_at = trade.timestamp;
+
+            // --- START: CORRECT UNIQUE HOLDER LOGIC ---
+
+            let current_pair = (trade.base_address.clone(), trade.maker.clone());
+
+            // We only increment the counter if:
+            // 1. The trader is NOT in the set of traders we know about from the database.
+            // 2. We have NOT already counted this trader for this token in this batch.
+            if !existing_traders.contains(&current_pair) {
+                // The .insert() returns true only the first time we see this pair in this batch.
+                if counted_in_this_batch.insert(current_pair) {
+                    entry.metrics.unique_holders += 1;
+                }
+            }
+
+            let trade_total_in_usd = trade.total_usd;
+
+            entry.metrics.total_volume_usd += trade_total_in_usd;
+            entry.metrics.ath_price_usd = entry.metrics.ath_price_usd.max(trade.price_usd);
+
+            if trade.trade_type == 0 {
+                // Buy
+                entry.metrics.total_buys += 1;
+            } else {
+                // Sell
+                entry.metrics.total_sells += 1;
+            }
+        }
+    }
+
+    async fn fetch_tokens_from_db(&self, keys: &[String]) -> Result<TokenCache> {
+        if keys.is_empty() {
+            return Ok(HashMap::new());
+        }
+        let query_str = "
+            SELECT
+                *
+            FROM tokens_latest
+            WHERE token_address IN ?
+        ";
+
+        let mut statics = HashMap::new();
+        for chunk in keys.chunks(*TOKEN_STATS_CHUNK_SIZE) {
+            let mut cursor = self
+                .db_client
+                .query(query_str)
+                .bind(chunk)
+                .fetch::<TokenStaticRow>()?;
+
+            while let Ok(Some(token)) = cursor.next().await {
+                statics.insert(token.token_address.clone(), token);
+            }
+        }
+
+        let metrics_map = self.fetch_token_metrics(keys).await?;
+        let mut tokens = HashMap::new();
+
+        for (address, token) in statics {
+            let metrics = metrics_map.get(&address).cloned();
+            tokens.insert(address.clone(), TokenEntry::new(token, metrics));
+        }
+
+        Ok(tokens)
+    }
+
+    async fn fetch_token_metrics(
+        &self,
+        keys: &[String],
+    ) -> Result<HashMap<String, TokenMetricsRow>> {
+        if keys.is_empty() {
+            return Ok(HashMap::new());
+        }
+
+        let query_str = "
+            SELECT
+                *
+            FROM token_metrics_latest
+            WHERE token_address IN ?
+            ORDER BY token_address, updated_at DESC
+            LIMIT 1 BY token_address
+        ";
+
+        let mut metrics = HashMap::new();
+
+        for chunk in keys.chunks(*TOKEN_STATS_CHUNK_SIZE) {
+            let mut cursor = self
+                .db_client
+                .query(query_str)
+                .bind(chunk)
+                .fetch::<TokenMetricsRow>()?;
+
+            while let Ok(Some(row)) = cursor.next().await {
+                metrics.insert(row.token_address.clone(), row);
+            }
+        }
+
+        Ok(metrics)
+    }
+
+    async fn fetch_token_metadata(&self, mint_address_str: &str) -> Result<(Metadata, Mint)> {
+        let mint_pubkey = Pubkey::from_str(mint_address_str)?;
+        let metadata_pubkey = Metadata::find_pda(&mint_pubkey).0;
+
+        let (mint_account_res, metadata_account_res) = future::join(
+            self.rpc_client.get_account(&mint_pubkey),
+            self.rpc_client.get_account(&metadata_pubkey),
+        )
+        .await;
+
+        let mint_account = mint_account_res?;
+        let metadata_account = metadata_account_res?;
+
+        let mint_data = Mint::unpack(&mint_account.data)?;
+        let metadata = Metadata::deserialize(&mut &metadata_account.data[..])?;
+
+        Ok((metadata, mint_data))
+    }
+
+    fn process_mint(&self, mint: &MintRow, tokens: &mut TokenCache) {
+        let is_new = !tokens.contains_key(&mint.mint_address);
+        let entry = tokens
+            .entry(mint.mint_address.clone())
+            .or_insert_with(|| TokenEntry::new(TokenStaticRow::new_from_mint(mint), None));
+        let token = &mut entry.token;
+
+        if is_new {
+            println!(
+                "[TokenAggregator] -> Created new token record for {} from MINT event.",
+                mint.mint_address
+            );
+        } else {
+            println!(
+                "[TokenAggregator] -> Enriched existing token record for {} with MINT event data.",
+                mint.mint_address
+            );
+            token.updated_at = mint.timestamp;
+            token.created_at = token.created_at.min(mint.timestamp);
+            token.decimals = mint.token_decimals;
+            token.launchpad = mint.protocol;
+            token.protocol = mint.protocol;
+            token.total_supply = mint.total_supply;
+            token.is_mutable = mint.is_mutable;
+            token.update_authority = mint.update_authority.clone();
+            token.mint_authority = mint.mint_authority.clone();
+            token.freeze_authority = mint.freeze_authority.clone();
+            if token.name.is_empty() {
+                token.name = mint.token_name.clone().unwrap_or_default();
+            }
+            if token.symbol.is_empty() {
+                token.symbol = mint.token_symbol.clone().unwrap_or_default();
+            }
+            if token.token_uri.is_empty() {
+                token.token_uri = mint.token_uri.clone().unwrap_or_default();
+            }
+            if token.creator_address.is_empty() {
+                token.creator_address = mint.creator_address.clone();
+            }
+            if !mint.pool_address.is_empty() && !token.pool_addresses.contains(&mint.pool_address) {
+                token.pool_addresses.push(mint.pool_address.clone());
+            }
+        }
+    }
+
+    fn process_migration(&self, migration: &MigrationRow, tokens: &mut TokenCache) {
+        if let Some(entry) = tokens.get_mut(&migration.mint_address) {
+            let token = &mut entry.token;
+            println!(
+                "[TokenAggregator] -> Updating protocol for token {} due to migration.",
+                migration.mint_address
+            );
+            token.updated_at = migration.timestamp;
+            token.protocol = migration.protocol;
+            if !token.pool_addresses.contains(&migration.pool_address) {
+                token.pool_addresses.push(migration.pool_address.clone());
+            }
+        }
+    }
+
+    async fn finalize_and_persist(&self, tokens: TokenCache) -> Result<()> {
+        if tokens.is_empty() {
+            return Ok(());
+        }
+
+        let mut updated_tokens = Vec::new();
+        let mut metric_rows = Vec::new();
+
+        for entry in tokens.into_values() {
+            if Self::metrics_has_activity(&entry.metrics) {
+                metric_rows.push(entry.metrics);
+            }
+            updated_tokens.push(entry.token);
+        }
+
+        insert_rows(
+            &self.db_client,
+            "tokens",
+            updated_tokens.clone(),
+            "Token Aggregator",
+            "tokens",
+        )
+        .await
+        .with_context(|| "Failed to persist token data to ClickHouse")?;
+
+        insert_rows(
+            &self.db_client,
+            "tokens_latest",
+            updated_tokens,
+            "Token Aggregator",
+            "tokens_latest",
+        )
+        .await
+        .with_context(|| "Failed to persist token snapshot data to ClickHouse")?;
+
+        insert_rows(
+            &self.db_client,
+            "token_metrics",
+            metric_rows.clone(),
+            "Token Aggregator",
+            "token_metrics",
+        )
+        .await
+        .with_context(|| "Failed to persist token metric history to ClickHouse")?;
+
+        insert_rows(
+            &self.db_client,
+            "token_metrics_latest",
+            metric_rows,
+            "Token Aggregator",
+            "token_metrics_latest",
+        )
+        .await
+        .with_context(|| "Failed to persist token metric snapshots to ClickHouse")?;
+
+        Ok(())
+    }
+
+    fn metrics_has_activity(metrics: &TokenMetricsRow) -> bool {
+        metrics.total_volume_usd > 0.0
+            || metrics.total_buys > 0
+            || metrics.total_sells > 0
+            || metrics.unique_holders > 0
+            || metrics.ath_price_usd > 0.0
+    }
+
+    async fn collect_events(
+        &mut self,
+        stream_key: &str,
+        group_name: &str,
+        consumer_name: &str,
+    ) -> Result<Vec<(String, EventPayload)>> {
+        let opts = StreamReadOptions::default()
+            .group(group_name, consumer_name)
+            .count(1000)
+            .block(2000);
+        let reply: StreamReadReply = self
+            .redis_conn
+            .xread_options(&[stream_key], &[">"], &opts)
+            .await?;
+        let mut events = Vec::new();
+        for stream_entry in reply.keys {
+            for message in stream_entry.ids {
+                if let Some(payload_value) = message.map.get("payload") {
+                    if let Ok(payload_bytes) = Vec::<u8>::from_redis_value(payload_value) {
+                        if let Ok(payload) = bincode::deserialize::<EventPayload>(&payload_bytes) {
+                            events.push((message.id.clone(), payload));
+                        }
+                    }
+                }
+            }
+        }
+        Ok(events)
+    }
+}

train.py

@@ -427,6 +427,7 @@ def main() -> None:
 
     # --- 7. Training Loop ---
     total_steps = 0
+    quality_loss_fn = nn.MSELoss()
     
     logger.info("***** Running training *****")
     logger.info(f"  Num examples = {len(dataset)}")
@@ -470,8 +471,12 @@ def main() -> None:
                 outputs = model(batch)
                 
                 preds = outputs["quantile_logits"]
+                quality_preds = outputs["quality_logits"]
                 labels = batch["labels"]
                 labels_mask = batch["labels_mask"]
+                if "quality_score" not in batch:
+                    raise RuntimeError("FATAL: quality_score missing from batch. Cannot train quality head.")
+                quality_targets = batch["quality_score"].to(accelerator.device, dtype=quality_preds.dtype)
                 if labels_mask is not None and labels_mask.sum().item() == 0:
                     token_addresses = batch.get('token_addresses', [])
                     t_cutoffs = batch.get('t_cutoffs', [])
@@ -482,11 +487,14 @@ def main() -> None:
                         token_addresses[0] if token_addresses else "unknown",
                         t_cutoffs[0] if t_cutoffs else "unknown",
                     )
-                
+
                 if labels_mask.sum() == 0:
-                    loss = torch.tensor(0.0, requires_grad=True, device=accelerator.device)
+                    return_loss = torch.tensor(0.0, requires_grad=True, device=accelerator.device)
                 else:
-                    loss = quantile_pinball_loss(preds, labels, labels_mask, quantiles)
+                    return_loss = quantile_pinball_loss(preds, labels, labels_mask, quantiles)
+
+                quality_loss = quality_loss_fn(quality_preds, quality_targets)
+                loss = return_loss + quality_loss
 
                 accelerator.backward(loss)
 
@@ -519,6 +527,8 @@ def main() -> None:
                      log_debug_batch_context(batch, logger, total_steps)
                 
                 current_loss = loss.item()
+                current_return_loss = return_loss.item()
+                current_quality_loss = quality_loss.item()
                 epoch_loss += current_loss
                 valid_batches += 1
                 
@@ -526,6 +536,8 @@ def main() -> None:
                     lr = scheduler.get_last_lr()[0]
                     log_payload = {
                         "train/loss": current_loss,
+                        "train/return_loss": current_return_loss,
+                        "train/quality_loss": current_quality_loss,
                         "train/learning_rate": lr,
                         "train/epoch": epoch + (step / len(dataloader))
                     }

train.sh

@@ -11,7 +11,7 @@ accelerate launch train.py \
   --tensorboard_dir runs/oracle \
   --checkpoint_dir checkpoints \
   --mixed_precision bf16 \
-  --max_seq_len 8192 \
+  --max_seq_len 4096 \
   --horizons_seconds 60 180 300 600 1800 3600 7200 \
   --quantiles 0.1 0.5 0.9 \
   --ohlc_stats_path ./data/ohlc_stats.npz \