Upload mixture/run_moe_inference.py with huggingface_hub

mixture/run_moe_inference.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python3
+"""Run inference using trained Task 1 (Modulation) and Task 2 (SNR/Mobility) MoE models.
+
+This script loads two separate MoE checkpoints and performs predictions on input spectrograms:
+  - Task 1 MoE: Predicts modulation scheme (QPSK, 16QAM, 64QAM, etc.)
+  - Task 2 MoE: Predicts joint SNR and mobility class
+
+Usage:
+    python -m mixture.run_moe_inference \\
+        --task1-checkpoint mixture/runs/task1_moe/moe_checkpoint.pth \\
+        --task2-checkpoint mixture/runs/task2_moe/moe_checkpoint.pth \\
+        --input spectrograms/city_1_losangeles/LTE/snr_0/pedestrian/QPSK/fft_512_overlap_256/specs_0000.pkl \\
+        --index 0
+
+Or run on a batch of samples:
+    python -m mixture.run_moe_inference \\
+        --task1-checkpoint mixture/runs/task1_moe/moe_checkpoint.pth \\
+        --task2-checkpoint mixture/runs/task2_moe/moe_checkpoint.pth \\
+        --input spectrograms/city_1_losangeles/LTE/snr_0/pedestrian/QPSK/fft_512_overlap_256/specs_0000.pkl \\
+        --batch-size 32
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+import torch
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+sys.path.append(str(REPO_ROOT))
+
+from mixture.train_embedding_router import MoEPredictor, load_all_samples  # type: ignore
+
+
+def load_spectrogram_sample(file_path: Path, index: Optional[int] = None) -> torch.Tensor:
+    """Load spectrogram(s) from pickle file.
+    
+    Args:
+        file_path: Path to pickle file containing spectrograms
+        index: If specified, return single sample at this index. Otherwise return all.
+    
+    Returns:
+        Tensor of shape [H, W] (single) or [N, H, W] (batch)
+    """
+    specs = load_all_samples(str(file_path))
+    
+    if index is not None:
+        if index < 0 or index >= specs.shape[0]:
+            raise IndexError(f"Index {index} out of range for file with {specs.shape[0]} samples")
+        return torch.from_numpy(specs[index]).float()
+    
+    return torch.from_numpy(specs).float()
+
+
+def format_prediction_output(result: dict, task_name: str) -> str:
+    """Format prediction result for console output."""
+    lines = [f"\n{task_name} Prediction:"]
+    lines.append("-" * 60)
+    
+    if "label" in result:
+        lines.append(f"  Predicted: {result['label']}")
+        lines.append(f"  Confidence: {result['confidence']:.4f}")
+    elif "labels" in result:
+        lines.append(f"  Batch size: {len(result['labels'])}")
+        lines.append(f"  Predictions: {result['labels'][:5]}{'...' if len(result['labels']) > 5 else ''}")
+    else:
+        lines.append(f"  Predicted class: {result['predicted_class']}")
+        lines.append(f"  Confidence: {result['confidence']:.4f}")
+    
+    if "routing" in result and result["routing"]:
+        lines.append("\n  Routing Weights:")
+        routing = result["routing"]
+        if isinstance(routing, list) and len(routing) > 0:
+            # Show routing for first sample in batch
+            if isinstance(routing[0], list):
+                routing = routing[0]
+        for expert_info in routing:
+            lines.append(f"    {expert_info['expert']:20s} ({expert_info['comm']:4s}): {expert_info['weight']:.4f}")
+    
+    return "\n".join(lines)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument(
+        "--task1-checkpoint",
+        type=Path,
+        required=True,
+        help="Path to Task 1 (modulation) MoE checkpoint",
+    )
+    parser.add_argument(
+        "--task2-checkpoint",
+        type=Path,
+        required=True,
+        help="Path to Task 2 (SNR/mobility) MoE checkpoint",
+    )
+    parser.add_argument(
+        "--input",
+        type=Path,
+        required=True,
+        help="Path to input spectrogram pickle file",
+    )
+    parser.add_argument(
+        "--index",
+        type=int,
+        default=None,
+        help="Index of sample to process (default: process all samples in file)",
+    )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=None,
+        help="If processing multiple samples, batch size for inference (default: process all at once)",
+    )
+    parser.add_argument(
+        "--show-probabilities",
+        action="store_true",
+        help="Show full class probability distributions",
+    )
+    parser.add_argument(
+        "--show-routing",
+        action="store_true",
+        help="Show expert routing weights",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=None,
+        help="Optional: save predictions to JSON file",
+    )
+    parser.add_argument(
+        "--device",
+        choices=["cuda", "cpu", "auto"],
+        default="auto",
+        help="Device to use for inference (default: auto-detect)",
+    )
+    args = parser.parse_args()
+    
+    # Set device
+    if args.device == "auto":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(args.device)
+    
+    print(f"[INFO] Using device: {device}")
+    
+    # Load MoE models
+    print(f"[INFO] Loading Task 1 MoE from {args.task1_checkpoint}")
+    task1_predictor = MoEPredictor.from_checkpoint(args.task1_checkpoint, device)
+    
+    print(f"[INFO] Loading Task 2 MoE from {args.task2_checkpoint}")
+    task2_predictor = MoEPredictor.from_checkpoint(args.task2_checkpoint, device)
+    
+    # Load input spectrogram(s)
+    print(f"[INFO] Loading spectrogram(s) from {args.input}")
+    spectrograms = load_spectrogram_sample(args.input, args.index)
+    
+    if spectrograms.dim() == 2:
+        print(f"[INFO] Processing single spectrogram of shape {tuple(spectrograms.shape)}")
+        num_samples = 1
+    else:
+        print(f"[INFO] Processing {spectrograms.shape[0]} spectrograms")
+        num_samples = spectrograms.shape[0]
+    
+    # Run inference
+    results = {"task1": [], "task2": []}
+    
+    if num_samples == 1 or args.batch_size is None:
+        # Single inference call
+        print("\n" + "="*60)
+        print("RUNNING INFERENCE")
+        print("="*60)
+        
+        task1_result = task1_predictor.predict(
+            spectrograms,
+            return_probabilities=args.show_probabilities,
+            return_routing=args.show_routing,
+        )
+        task2_result = task2_predictor.predict(
+            spectrograms,
+            return_probabilities=args.show_probabilities,
+            return_routing=args.show_routing,
+        )
+        
+        results["task1"] = [task1_result] if num_samples == 1 else task1_result
+        results["task2"] = [task2_result] if num_samples == 1 else task2_result
+        
+        # Print results
+        print(format_prediction_output(task1_result, "Task 1 (Modulation)"))
+        print(format_prediction_output(task2_result, "Task 2 (SNR/Mobility)"))
+        
+    else:
+        # Batch processing
+        print("\n" + "="*60)
+        print(f"RUNNING BATCH INFERENCE ({args.batch_size} samples at a time)")
+        print("="*60)
+        
+        num_batches = (num_samples + args.batch_size - 1) // args.batch_size
+        
+        for batch_idx in range(num_batches):
+            start_idx = batch_idx * args.batch_size
+            end_idx = min(start_idx + args.batch_size, num_samples)
+            batch_specs = spectrograms[start_idx:end_idx]
+            
+            print(f"\n[Batch {batch_idx+1}/{num_batches}] Processing samples {start_idx} to {end_idx-1}")
+            
+            task1_batch_result = task1_predictor.predict(
+                batch_specs,
+                return_probabilities=args.show_probabilities,
+                return_routing=args.show_routing,
+            )
+            task2_batch_result = task2_predictor.predict(
+                batch_specs,
+                return_probabilities=args.show_probabilities,
+                return_routing=args.show_routing,
+            )
+            
+            results["task1"].extend(
+                [task1_batch_result] if isinstance(task1_batch_result.get("predicted_class"), int)
+                else [{"predicted_class": task1_batch_result["predicted_class"][i],
+                       "label": task1_batch_result.get("labels", [None])[i],
+                       "confidence": task1_batch_result.get("confidence")[i] if isinstance(task1_batch_result.get("confidence"), list) else task1_batch_result.get("confidence")}
+                      for i in range(len(task1_batch_result.get("labels", task1_batch_result.get("predicted_class", []))))]
+            )
+            results["task2"].extend(
+                [task2_batch_result] if isinstance(task2_batch_result.get("predicted_class"), int)
+                else [{"predicted_class": task2_batch_result["predicted_class"][i],
+                       "label": task2_batch_result.get("labels", [None])[i],
+                       "confidence": task2_batch_result.get("confidence")[i] if isinstance(task2_batch_result.get("confidence"), list) else task2_batch_result.get("confidence")}
+                      for i in range(len(task2_batch_result.get("labels", task2_batch_result.get("predicted_class", []))))]
+            )
+        
+        # Print summary
+        print("\n" + "="*60)
+        print("INFERENCE SUMMARY")
+        print("="*60)
+        print(f"Total samples processed: {num_samples}")
+        
+        if results["task1"]:
+            task1_labels = [r.get("label", "Unknown") for r in results["task1"]]
+            print(f"\nTask 1 (Modulation) predictions:")
+            unique_labels = set(task1_labels)
+            for label in sorted(unique_labels):
+                count = task1_labels.count(label)
+                print(f"  {label}: {count} samples ({count/num_samples*100:.1f}%)")
+        
+        if results["task2"]:
+            task2_labels = [r.get("label", "Unknown") for r in results["task2"]]
+            print(f"\nTask 2 (SNR/Mobility) predictions:")
+            unique_labels = set(task2_labels)
+            for label in sorted(unique_labels):
+                count = task2_labels.count(label)
+                print(f"  {label}: {count} samples ({count/num_samples*100:.1f}%)")
+    
+    # Save results to file if requested
+    if args.output:
+        output_path = args.output.expanduser().resolve()
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        # Convert tensors to Python types for JSON serialization
+        def jsonify(obj):
+            if isinstance(obj, dict):
+                return {k: jsonify(v) for k, v in obj.items()}
+            elif isinstance(obj, (list, tuple)):
+                return [jsonify(x) for x in obj]
+            elif isinstance(obj, (torch.Tensor, np.ndarray)):
+                return obj.tolist()
+            elif isinstance(obj, (np.integer, np.floating)):
+                return obj.item()
+            return obj
+        
+        output_data = {
+            "input_file": str(args.input),
+            "num_samples": num_samples,
+            "task1_predictions": jsonify(results["task1"]),
+            "task2_predictions": jsonify(results["task2"]),
+        }
+        
+        with output_path.open("w", encoding="utf-8") as f:
+            json.dump(output_data, f, indent=2)
+        
+        print(f"\n[INFO] Results saved to {output_path}")
+    
+    print("\n" + "="*60)
+    print("INFERENCE COMPLETE")
+    print("="*60 + "\n")
+
+
+if __name__ == "__main__":
+    main()
+