Add flexible safetensors file loading

- Add --weights argument to specify custom .safetensors file
- Auto-discover any .safetensors file if model.safetensors not found
- Update discover_circuits to accept any .safetensors file

prune.py

@@ -495,9 +495,10 @@ class ComputationGraph:
 class AdaptiveCircuit:
     """Adaptive threshold circuit with automatic batched evaluation."""
 
-    def __init__(self, path: Path, device: str = 'cuda'):
+    def __init__(self, path: Path, device: str = 'cuda', weights_file: str = None):
         self.path = Path(path)
         self.device = device
+        self.weights_file = weights_file
         self.spec = self._load_spec()
         self.weights = self._load_weights()
         self.weight_keys = list(self.weights.keys())
@@ -655,7 +656,17 @@ class AdaptiveCircuit:
         )
 
     def _load_weights(self) -> Dict[str, torch.Tensor]:
-        w = load_file(str(self.path / 'model.safetensors'))
+        if self.weights_file:
+            safetensors_file = self.path / self.weights_file
+        else:
+            safetensors_file = self.path / 'model.safetensors'
+            if not safetensors_file.exists():
+                candidates = list(self.path.glob('*.safetensors'))
+                if candidates:
+                    safetensors_file = candidates[0]
+                else:
+                    raise FileNotFoundError(f"No .safetensors file found in {self.path}")
+        w = load_file(str(safetensors_file))
         return {k: v.float().to(self.device) for k, v in w.items()}
 
     def _build_tests(self) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -3010,7 +3021,7 @@ def discover_circuits(base: Path = CIRCUITS_PATH) -> List[CircuitSpec]:
     """Find all circuits."""
     circuits = []
     for d in base.iterdir():
-        if d.is_dir() and (d / 'config.json').exists() and (d / 'model.safetensors').exists():
+        if d.is_dir() and (d / 'config.json').exists() and list(d.glob('*.safetensors')):
             try:
                 with open(d / 'config.json') as f:
                     cfg = json.load(f)
@@ -3105,6 +3116,7 @@ def run_pipeline(circuit: AdaptiveCircuit, evaluator: BatchedEvaluator, cfg: Con
 def main():
     parser = argparse.ArgumentParser(description='Threshold Circuit Pruner v4')
     parser.add_argument('circuit', nargs='?', help='Circuit name')
+    parser.add_argument('--weights', type=str, help='Specific .safetensors file to load (default: model.safetensors or first found)')
     parser.add_argument('--list', action='store_true')
     parser.add_argument('--all', action='store_true')
     parser.add_argument('--max-inputs', type=int, default=10)
@@ -3198,7 +3210,7 @@ def main():
             print(f"Circuit not found: {args.circuit}")
             return
 
-        circuit = AdaptiveCircuit(path, cfg.device)
+        circuit = AdaptiveCircuit(path, cfg.device, args.weights)
 
         if args.pipeline:
             stages = [s.strip() for s in args.pipeline.split(',')]