Randomize candidate dataset order

dataset/README.md

@@ -103,7 +103,7 @@ Verify:
 
 ```bash
 modal volume ls gazet-data
-# should show: overture/, overture_normalized/, natural_earth_geoparquet/, natural_earth_normalized/
+# should show: overture_normalized/, natural_earth_normalized/
 ```
 
 Skip this step on subsequent runs — the volume persists across runs.

dataset/scripts/export_training_data.py

@@ -20,6 +20,7 @@ Output layout (all paths relative to dataset/):
     output/runs/{run_name}/stats.json
 """
 
+import copy
 import json
 import random
 import sys
@@ -178,6 +179,23 @@ def _format_sql(sql: str) -> str:
     ).strip()
 
 
+def _shuffle_candidates_for_export(sample: Dict[str, Any]) -> Dict[str, Any]:
+    """Return a copy of sample with candidate row order shuffled deterministically.
+
+    Candidate IDs remain unchanged; only row order changes. This removes the
+    positional shortcut where the true anchor often appears first in the raw
+    synthetic samples, while keeping selected_candidates valid.
+    """
+    shuffled = copy.deepcopy(sample)
+    candidates = shuffled.get("candidates", [])
+    if len(candidates) <= 1:
+        return shuffled
+
+    rng = random.Random(shuffled.get("id", "sample"))
+    rng.shuffle(candidates)
+    return shuffled
+
+
 def sample_to_sql_pair(sample: Dict[str, Any]) -> Optional[Dict]:
     """Convert a raw sample to a conversational prompt-completion pair for SQL generation."""
     sql = sample.get("target", {}).get("sql", "").strip()
@@ -360,7 +378,9 @@ def main(config_path: Optional[Path] = None) -> None:
     # Load
     print("\nLoading validated samples...")
     samples = load_samples(validated_file)
+    samples = [_shuffle_candidates_for_export(s) for s in samples]
     print(f"  {len(samples):,} samples loaded")
+    print("  Candidate row order shuffled deterministically for export")
 
     # Split once, reuse for both tasks
     print("\nSplitting 80 / 10 / 10 (stratified by task family)...")

dataset/scripts/generate_samples.py

@@ -416,6 +416,35 @@ def _merge_candidate_lists(
     return merged
 
 
+def _dedupe_country_candidates(
+    candidates: List[Candidate],
+    max_total: Optional[int] = None,
+) -> List[Candidate]:
+    """Deduplicate country candidates by country code, preserving first match.
+
+    This is useful for templates whose SQL uses ``country IN (...)`` rather
+    than candidate IDs. Overture can contain multiple country-level rows for
+    the same ISO code, which weakens grounding if they all remain in the list.
+    """
+    deduped: List[Candidate] = []
+    seen_keys: set[tuple[str, str]] = set()
+    for cand in candidates:
+        if cand.subtype == "country" and cand.country:
+            key = ("country", cand.country)
+        else:
+            key = ("id", cand.id)
+        if key in seen_keys:
+            continue
+        deduped.append(cand)
+        seen_keys.add(key)
+        if max_total is not None and len(deduped) >= max_total:
+            break
+
+    for i, cand in enumerate(deduped, 1):
+        cand.candidate_id = f"c{i}"
+    return deduped
+
+
 def build_candidate_list(
     con: duckdb.DuckDBPyConnection,
     anchor_id: str,
@@ -1008,19 +1037,29 @@ def generate_template_based_sample(
             )
 
         elif template.template_id in ("contain_multi_01", "contain_multi_02", "contain_multi_03"):
-            # country IN clause — 2 or 3 anchors, each contributes its country code
+            # country IN clause — 2 or 3 anchors, each contributes its country code.
+            # Sample unique countries so the query actually teaches a multi-country
+            # pattern rather than repeating the same ISO code multiple times.
             num_a = 3 if template.template_id == "contain_multi_02" else 2
+            country_inventory = tables['divisions_area_inventory']
+            country_inventory = country_inventory[
+                (country_inventory['subtype'] == 'country')
+                & country_inventory['country'].notna()
+            ].drop_duplicates(subset=['country'])
+            if len(country_inventory) < num_a:
+                return None
+
+            sampled = country_inventory.sample(n=num_a, replace=False)
             anchors = [
-                sample_random_entity(
-                    con,
-                    tables['divisions_area_inventory'],
-                    'divisions_area',
-                    subtypes={'country'},
-                )
-                for _ in range(num_a)
+                {
+                    'id': row['id'],
+                    'name': row['name'],
+                    'subtype': row.get('subtype'),
+                    'country': row.get('country'),
+                    'source': 'divisions_area',
+                }
+                for _, row in sampled.iterrows()
             ]
-            if any(a is None for a in anchors):
-                return None
 
             countries = [a.get('country') or 'US' for a in anchors]
             target_subtype = template.target_subtype or 'region'
@@ -1039,7 +1078,10 @@ def generate_template_based_sample(
                                      num_candidates=per_anchor, difficulty="medium")
                 for a in anchors
             ]
-            candidates = _merge_candidate_lists(*cands, max_total=num_a * per_anchor)
+            candidates = _dedupe_country_candidates(
+                _merge_candidate_lists(*cands, max_total=num_a * per_anchor),
+                max_total=num_a * per_anchor,
+            )
 
             q_kwargs = dict(target_subtype=target_subtype)
             for i, a in enumerate(anchors, 1):
@@ -1629,7 +1671,12 @@ def generate_template_based_sample(
         template.anchor_source == "mixed" and template.num_anchors == 2
     )
 
-    if template.family in _multi_anchor_families and template.num_anchors >= 2:
+    if _is_mixed_two_anchor:
+        # partial_05 / diff_02: anchor (division) + clip_feature (natural_earth)
+        mixed_ids = {anchor.get("id", ""), clip_feature.get("id", "")}
+        selected_candidate_ids = [c.candidate_id for c in candidates if c.id in mixed_ids]
+
+    elif template.family in _multi_anchor_families and template.num_anchors >= 2:
         anchor_ids: set = set()
         for var in ("anchor1", "anchor2", "anchor3"):
             obj = locals().get(var)
@@ -1641,11 +1688,6 @@ def generate_template_based_sample(
                     anchor_ids.add(a.get("id", ""))
         selected_candidate_ids = [c.candidate_id for c in candidates if c.id in anchor_ids]
 
-    elif _is_mixed_two_anchor:
-        # partial_05 / diff_02: anchor (division) + clip_feature (natural_earth)
-        mixed_ids = {anchor.get("id", ""), clip_feature.get("id", "")}
-        selected_candidate_ids = [c.candidate_id for c in candidates if c.id in mixed_ids]
-
     else:
         anchor_id_to_find = (
             anchor.get('anchor_id')