feat: add training progress prints to from_training_data()

router.py

@@ -137,8 +137,12 @@ class R2Router:
         with open(os.path.join(path, "training_data", "labels.json")) as f:
             labels = json.load(f)
 
+        print(f"Training KNN (k={k}) on {len(X_train)} samples...")
+
         quality_knns = {}
         token_knns = {}
+        n_quality = 0
+        n_token = 0
 
         for model_name, model_labels in labels.items():
             quality_knns[model_name] = {}
@@ -154,6 +158,7 @@ class R2Router:
                 )
                 knn.fit(X_train[valid], acc[valid])
                 quality_knns[model_name][budget_name] = knn
+                n_quality += 1
 
             if "concise" in model_labels and "output_tokens" in model_labels["concise"]:
                 tok = np.array([x if x is not None else np.nan for x in model_labels["concise"]["output_tokens"]])
@@ -166,6 +171,9 @@ class R2Router:
                     )
                     tknn.fit(X_train[valid], tok[valid])
                     token_knns[model_name] = tknn
+                    n_token += 1
+
+        print(f"Trained {n_quality} quality KNNs + {n_token} token KNNs for {len(quality_knns)} models.")
 
         model_prices = {
             mn: cfg["output_price_per_million"]