Makefile

@@ -16,7 +16,7 @@ eval: eval.py training.csv
 
 train_frozen_encoder: train.py training.csv
 	@echo "Training coordinate head with frozen encoder..."
-	@bash -c 'source .venv/bin/activate && python train.py --output-path output_frozen_encoder --freeze-encoder --epochs 10'
+	@bash -c 'source .venv/bin/activate && python train.py --output-path output_frozen_encoder --freeze-encoder --epochs 50'
 
 eval_frozen_encoder: eval.py training.csv
 	@echo "Evaluating frozen-encoder coordinate regressor..."
@@ -24,7 +24,7 @@ eval_frozen_encoder: eval.py training.csv
 
 train_frozen_layers: train.py training.csv
 	@echo "Training coordinate regressor with first transformer layers frozen..."
-	@bash -c 'source .venv/bin/activate && python train.py --output-path output_frozen_layers --freeze-transformer-layers 4 --epochs 10'
+	@bash -c 'source .venv/bin/activate && python train.py --output-path output_frozen_layers --freeze-transformer-layers 4 --epochs 50'
 
 eval_frozen_layers: eval.py training.csv
 	@echo "Evaluating frozen-layer coordinate regressor..."

data_utils.py

@@ -1,27 +0,0 @@
-import numpy as np
-from sklearn.model_selection import GroupShuffleSplit, train_test_split
-
-DEFAULT_TEST_SIZE = 0.2
-DEFAULT_SEED = 1992
-
-
-def split_indices(data, test_size=DEFAULT_TEST_SIZE, seed=DEFAULT_SEED):
-    """Return (train_indices, val_indices) arrays.
-
-    Uses a group-aware split on the ``intersection`` column when present so
-    that every row from the same intersection lands in the same partition.
-    Falls back to a plain random split otherwise.
-    """
-    indices = np.arange(len(data))
-    if "intersection" in data.columns:
-        splitter = GroupShuffleSplit(
-            n_splits=1, test_size=test_size, random_state=seed
-        )
-        train_idx, val_idx = next(
-            splitter.split(indices, groups=data["intersection"])
-        )
-    else:
-        train_idx, val_idx = train_test_split(
-            indices, test_size=test_size, random_state=seed
-        )
-    return train_idx, val_idx

eval.py

@@ -11,7 +11,6 @@ import pandas as pd  # noqa: E402
 import torch  # noqa: E402
 from sentence_transformers import SentenceTransformer  # noqa: E402
 
-from data_utils import split_indices  # noqa: E402
 from train import CoordinateRegressor, haversine_km  # noqa: E402
 
 matplotlib.use("Agg")
@@ -35,18 +34,6 @@ def parse_args():
         help="Device to use for evaluation. Defaults to `cuda`.",
     )
     parser.add_argument("--batch-size", type=int, default=64)
-    parser.add_argument(
-        "--seed",
-        type=int,
-        default=1992,
-        help="Must match the seed used during training for a consistent split.",
-    )
-    parser.add_argument(
-        "--test-size",
-        type=float,
-        default=0.2,
-        help="Must match the test-size used during training.",
-    )
     return parser.parse_args()
 
 
@@ -231,10 +218,7 @@ def main():
         coord_std,
     ) = load_model(args.model_path, device)
 
-    data = pd.read_csv(args.data_file).dropna(
+    data = pd.read_csv(args.data_file).dropna(subset=["text", "latitude", "longitude"])
-        subset=["text", "latitude", "longitude"]
-    )
-    data = data.reset_index(drop=True)
     texts = data["text"].astype(str).tolist()
     predicted = predict(
         trained_encoder,
@@ -264,15 +248,6 @@ def main():
         results["initial_predicted_latitude"] = initial_predicted[:, 0]
         results["initial_predicted_longitude"] = initial_predicted[:, 1]
     results["error_km"] = errors
-
-    # Reproduce the same train/val split used during training.
-    train_indices, val_indices = split_indices(
-        data, test_size=args.test_size, seed=args.seed
-    )
-    val_mask = np.zeros(len(results), dtype=bool)
-    val_mask[val_indices] = True
-    results["split"] = np.where(val_mask, "val", "train")
-
     results.to_csv(args.output_file, index=False)
     make_prediction_plot(results, args.plot_file)
     make_predicted_vs_actual_plot(results, args.scatter_plot_file)
@@ -280,48 +255,9 @@ def main():
     print(f"Wrote predictions to {args.output_file}")
     print(f"Wrote plot to {args.plot_file}")
     print(f"Wrote scatter plot to {args.scatter_plot_file}")
-
+    print(f"mean_error_km={errors.mean():.3f}")
-    def _fmt(errs):
+    print(f"median_error_km={np.median(errors):.3f}")
-        return (
+    print(f"p90_error_km={np.percentile(errors, 90):.3f}")
-            f"mean={np.mean(errs):.3f}  "
-            f"median={np.median(errs):.3f}  "
-            f"p90={np.percentile(errs, 90):.3f}"
-        )
-
-    def _intersection_agg_errors(df):
-        """Average per-bag predictions to one coordinate per intersection."""
-        agg = (
-            df.groupby("intersection")
-            .agg(
-                pred_lat=("predicted_latitude", "mean"),
-                pred_lon=("predicted_longitude", "mean"),
-                true_lat=("latitude", "first"),
-                true_lon=("longitude", "first"),
-            )
-            .reset_index()
-        )
-        return haversine_km(
-            agg[["pred_lat", "pred_lon"]].to_numpy(dtype=np.float32),
-            agg[["true_lat", "true_lon"]].to_numpy(dtype=np.float32),
-        )
-
-    val_df = results[results["split"] == "val"]
-    train_df = results[results["split"] == "train"]
-
-    val_bag_errors = val_df["error_km"].to_numpy()
-    train_bag_errors = train_df["error_km"].to_numpy()
-    val_agg_errors = _intersection_agg_errors(val_df)
-    train_agg_errors = _intersection_agg_errors(train_df)
-
-    n_val = val_df["intersection"].nunique() if "intersection" in val_df.columns else "?"
-    n_train = train_df["intersection"].nunique() if "intersection" in train_df.columns else "?"
-
-    print()
-    print(f"[all bags  ({len(errors):>6} rows)]  {_fmt(errors)}")
-    print(f"[train bags({len(train_bag_errors):>6} rows)]  {_fmt(train_bag_errors)}")
-    print(f"[val bags  ({len(val_bag_errors):>6} rows)]  {_fmt(val_bag_errors)}")
-    print(f"[train agg ({n_train:>3} intersections)]  {_fmt(train_agg_errors)}")
-    print(f"[val agg   ({n_val:>3} intersections)]  {_fmt(val_agg_errors)}  ← generalization")
 
 
 if __name__ == "__main__":

prepare_training_data.py

@@ -1,10 +1,8 @@
 import argparse
-import random
 
 import pandas as pd
 
-BASE_COLUMNS = ["intersection", "text_on_sign_exact", "latitude", "longitude"]
+INPUT_COLUMNS = ["intersection", "text_on_sign_exact", "latitude", "longitude"]
-OPTIONAL_COLUMNS = ["code_type"]
 EXCLUDED_INTERSECTIONS = {"56th-pena"}
 
 
@@ -33,11 +31,8 @@ def parse_args():
     parser.add_argument(
         "--bag-size",
         type=int,
-        default=8,
+        default=5,
-        help=(
+        help="Number of sign texts to sample for each training row.",
-            "Total number of sign texts per training row. Street signs fill "
-            "slots first; remaining slots are randomly sampled from all signs."
-        ),
     )
     parser.add_argument(
         "--samples-per-intersection",
@@ -55,13 +50,12 @@ def parse_args():
 
 def load_raw_data(path):
     data = pd.read_csv(path)
-    missing = sorted(set(BASE_COLUMNS) - set(data.columns))
+    missing = sorted(set(INPUT_COLUMNS) - set(data.columns))
     if missing:
         raise ValueError(f"Missing required columns: {missing}")
 
-    keep = BASE_COLUMNS + [c for c in OPTIONAL_COLUMNS if c in data.columns]
+    data = data[INPUT_COLUMNS].copy()
-    data = data[keep].copy()
+    data = data.dropna(subset=INPUT_COLUMNS)
-    data = data.dropna(subset=BASE_COLUMNS)
     data["text_on_sign_exact"] = data["text_on_sign_exact"].astype(str).str.strip()
     data = data[data["text_on_sign_exact"] != ""]
     data = data[~data["intersection"].isin(EXCLUDED_INTERSECTIONS)]
@@ -74,48 +68,28 @@ def make_bootstrap_data(data, bag_size, samples_per_intersection, seed, separato
     if samples_per_intersection < 1:
         raise ValueError("--samples-per-intersection must be at least 1")
 
-    has_code_type = "code_type" in data.columns
     rows = []
     grouped = data.groupby("intersection", sort=True)
     for intersection, group in grouped:
+        texts = group["text_on_sign_exact"].tolist()
         latitude = group["latitude"].mean()
         longitude = group["longitude"].mean()
-        all_texts = group["text_on_sign_exact"].tolist()
-
-        # Street signs are the most geographically discriminative; always
-        # include them first up to bag_size, then fill remaining slots randomly.
-        if has_code_type:
-            street_texts = group[group["code_type"] == "street_sign"][
-                "text_on_sign_exact"
-            ].tolist()
-        else:
-            street_texts = []
 
         for sample_id in range(samples_per_intersection):
-            rng_state = seed + len(rows)
+            sampled = group.sample(
-
+                n=bag_size,
-            guaranteed = street_texts[:bag_size]
+                replace=True,
-            remaining = bag_size - len(guaranteed)
+                random_state=seed + len(rows),
-            if remaining > 0:
+            )
-                filler = group.sample(
-                    n=remaining, replace=True, random_state=rng_state
-                )["text_on_sign_exact"].tolist()
-            else:
-                filler = []
-
-            bag = guaranteed + filler
-            # Shuffle so position in the sequence doesn't encode sign type.
-            random.Random(rng_state).shuffle(bag)
-
             rows.append(
                 {
                     "intersection": intersection,
                     "sample_id": sample_id,
-                    "text": separator.join(bag),
+                    "text": separator.join(sampled["text_on_sign_exact"].tolist()),
                     "latitude": latitude,
                     "longitude": longitude,
-                    "unique_sign_count": len(set(all_texts)),
+                    "unique_sign_count": len(set(texts)),
-                    "raw_sign_count": len(all_texts),
+                    "raw_sign_count": len(texts),
                 }
             )
 

train.py

@@ -8,7 +8,7 @@ import numpy as np
 import pandas as pd
 import torch
 from sentence_transformers import LoggingHandler, SentenceTransformer
-from data_utils import split_indices
+from sklearn.model_selection import GroupShuffleSplit, train_test_split
 from torch import nn
 from torch.utils.data import DataLoader, Dataset
 
@@ -91,9 +91,9 @@ def parse_args():
             "epoch). Validation still runs every epoch."
         ),
     )
-    parser.add_argument("--learning-rate", type=float, default=1e-4)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
-    parser.add_argument("--head-learning-rate", type=float, default=5e-2)
+    parser.add_argument("--head-learning-rate", type=float, default=1e-3)
-    parser.add_argument("--weight-decay", type=float, default=0.001)
+    parser.add_argument("--weight-decay", type=float, default=0.01)
     parser.add_argument("--test-size", type=float, default=0.2)
     parser.add_argument("--hidden-dim", type=int, default=256)
     parser.add_argument("--dropout", type=float, default=0.1)
@@ -433,8 +433,17 @@ def main():
     coordinates = data[["latitude", "longitude"]].to_numpy(dtype=np.float32)
     normalized_coordinates, coord_mean, coord_std = normalize_coordinates(coordinates)
 
-    train_indices, val_indices = split_indices(
+    indices = np.arange(len(data))
-        data, test_size=args.test_size, seed=args.seed
+    if "intersection" in data.columns:
+        splitter = GroupShuffleSplit(
+            n_splits=1, test_size=args.test_size, random_state=args.seed
+        )
+        train_indices, val_indices = next(
+            splitter.split(indices, groups=data["intersection"])
+        )
+    else:
+        train_indices, val_indices = train_test_split(
+            indices, test_size=args.test_size, random_state=args.seed
         )
 
     train_dataset = SignCoordinateDataset(
@@ -506,10 +515,7 @@ def main():
         )
 
     optimizer = make_optimizer(encoder, head, args)
-    loss_fn = nn.HuberLoss()
+    loss_fn = nn.MSELoss()
-    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
-        optimizer, mode="min", patience=5, factor=0.5, min_lr=1e-7
-    )
     best_val_loss = float("inf")
     best_states = None
     pending_save = False
@@ -546,12 +552,9 @@ def main():
                 coord_std,
                 device,
             )
-        scheduler.step(val_loss)
-        current_lr = optimizer.param_groups[-1]["lr"]
         print(
             f"epoch={epoch} train_loss={train_loss:.6f} "
             f"val_loss={val_loss:.6f} val_error_km={val_error_km:.3f}"
-            f"lr={current_lr:.2e}"
         )
 
         if val_loss < best_val_loss: