train.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "sentence-transformers>=3.0.0",
#     "torch",
#     "transformers>=4.40.0",
#     "numpy",
#     "einops",
#     "datasets",
#     "accelerate>=1.1.0",
# ]
#
# [tool.uv.sources]
# torch = [
#     { index = "pytorch-cu124", marker = "sys_platform == 'linux'" },
# ]
#
# [[tool.uv.index]]
# name = "pytorch-cu124"
# url = "https://download.pytorch.org/whl/cu124"
# explicit = true
# ///
"""
Self-contained automotive embedding fine-tuning + evaluation.

Fine-tunes nomic-embed-text-v1.5 to match repair concerns → PCAdb subcategories.
Outputs Hit@1/5/10 on a 55-concern benchmark.

Run locally:  uv run train.py
Run on HF:    hf jobs uv run --flavor a100-large --timeout 30m --namespace <ns> train.py
"""

import json
import os
import random
import sys
import tempfile
from pathlib import Path

import numpy as np
from sentence_transformers import InputExample, SentenceTransformer, losses
from torch.utils.data import DataLoader

random.seed(42)
np.random.seed(42)

# ============================================================================
# BENCHMARK — DO NOT MODIFY
# 55 concerns (49 real + 1 garbage). This is the fixed evaluation set.
# ============================================================================

BENCHMARK = [
    {"concern": "Replace front brake pads and rotors", "expected_subcategories": ["Brake > Disc Pads and Brake Shoes", "Brake > Drums and Rotors"]},
    {"concern": "Brake pads worn to backing", "expected_subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    {"concern": "Replace front shocks and struts", "expected_subcategories": ["Suspension > Shocks and Struts"]},
    {"concern": "Control arm bushing worn", "expected_subcategories": ["Suspension > Suspension, Springs and Related Components"]},
    {"concern": "Timing chain replacement needed", "expected_subcategories": ["Engine > Valve Train Components"]},
    {"concern": "Serpentine belt cracking", "expected_subcategories": ["Belts and Cooling > Accessory Drive Belt System Components"]},
    {"concern": "Battery dead needs replacement", "expected_subcategories": ["Electrical, Charging and Starting > Battery and Related Components"]},
    {"concern": "Alternator failing low voltage", "expected_subcategories": ["Electrical, Charging and Starting > Alternator / Generator and Related Components"]},
    {"concern": "Fuel pump making noise", "expected_subcategories": ["Air and Fuel Delivery > Fuel Pumps and Related Components"]},
    {"concern": "A/C compressor not engaging", "expected_subcategories": ["HVAC > A/C Clutch and Compressor"]},
    {"concern": "Power steering pump leak", "expected_subcategories": ["Steering > Power Steering Hoses, Pumps, and Related Components"]},
    {"concern": "Spark plugs replacement due", "expected_subcategories": ["Ignition > Secondary Ignition"]},
    {"concern": "Wiper blades streaking", "expected_subcategories": ["Wiper and Washer > Wiper Arm, Blade, and Related Components"]},
    {"concern": "Wheel bearing noise from front", "expected_subcategories": ["Driveline and Axles > Wheel Bearings, Seals, and Related Components"]},
    {"concern": "Synthetic Oil Service due", "expected_subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "Power steering flush needed", "expected_subcategories": ["Steering > Power Steering Hoses, Pumps, and Related Components", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "Differential service recommended", "expected_subcategories": ["Driveline and Axles > Differential", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "Transfer case fluid change", "expected_subcategories": ["Transfer Case > Transfer Case Components", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "BRAKE PAD WORN", "expected_subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    {"concern": "AC recharge needed blowing warm", "expected_subcategories": ["HVAC > A/C Clutch and Compressor"]},
    {"concern": "r&r radiator", "expected_subcategories": ["Belts and Cooling > Radiators, Coolers and Related Components"]},
    {"concern": "Vehicle pulls right", "expected_subcategories": ["Suspension > Alignment"]},
    {"concern": "CEL on rough running", "expected_subcategories": ["Ignition > Secondary Ignition"]},
    {"concern": "Mailer oil change", "expected_subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "ATF flush", "expected_subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"concern": "Brak pads worn", "expected_subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    {"concern": "upper mount plate", "expected_subcategories": ["Suspension > Shocks and Struts"]},
    {"concern": "mirror assembly", "expected_subcategories": ["Body > Mirrors"]},
    {"concern": "grinding noise when braking", "expected_subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    {"concern": "car won't start clicking", "expected_subcategories": ["Electrical, Charging and Starting > Starter and Related Components"]},
    {"concern": "clunking over bumps", "expected_subcategories": ["Suspension > Shocks and Struts"]},
    {"concern": "overheating", "expected_subcategories": ["Belts and Cooling > Radiators, Coolers and Related Components"]},
    {"concern": "a/c blowing warm", "expected_subcategories": ["HVAC > A/C Clutch and Compressor"]},
    {"concern": "rotten egg smell", "expected_subcategories": ["Exhaust > Catalytic Converter"]},
    {"concern": "clicking when turning", "expected_subcategories": ["Driveline and Axles > CV Components"]},
    {"concern": "whining when turning", "expected_subcategories": ["Steering > Power Steering Hoses, Pumps, and Related Components"]},
    {"concern": "humming from wheel area", "expected_subcategories": ["Driveline and Axles > Wheel Bearings, Seals, and Related Components"]},
    {"concern": "squealing belt noise", "expected_subcategories": ["Belts and Cooling > Accessory Drive Belt System Components"]},
    {"concern": "oil leak from top of engine", "expected_subcategories": ["Engine > Gaskets and Sealing Systems"]},
    {"concern": "loud exhaust", "expected_subcategories": ["Exhaust > Mufflers"]},
    {"concern": "Replace catalytic converter", "expected_subcategories": ["Exhaust > Catalytic Converter"]},
    {"concern": "CV axle boot torn", "expected_subcategories": ["Driveline and Axles > CV Components", "Driveline and Axles > CV Boots and Related Components"]},
    {"concern": "Cabin air filter replacement", "expected_subcategories": ["HVAC > Filters"]},
    {"concern": "Blower motor not working", "expected_subcategories": ["HVAC > Motors, Core, Case and Related Components"]},
    {"concern": "EGR valve stuck", "expected_subcategories": ["Emission Control > EGR and Related Components"]},
    {"concern": "TPMS light on", "expected_subcategories": ["Tire and Wheel > Sensors"]},
    {"concern": "New tires needed", "expected_subcategories": ["Tire and Wheel > Tire"]},
    {"concern": "Turbo boost leak", "expected_subcategories": ["Air and Fuel Delivery > Turbocharger, Supercharger and Ram Air"]},
    {"concern": "customer wants alignment check", "expected_subcategories": ["Suspension > Alignment"]},
    {"concern": "Complete garbage text 12345 xyz", "expected_subcategories": []},
]

# ============================================================================
# === EXPERIMENT ZONE === (modify freely below this line)
# ============================================================================

# -- Hyperparameters --
BASE_MODEL = "nomic-ai/nomic-embed-text-v1.5"
EPOCHS = 15
BATCH_SIZE = 64
LEARNING_RATE = 1.5e-4
WARMUP_RATIO = 0.1
MATRYOSHKA_DIMS = []  # Disabled — dilutes gradient across dims
GARBAGE_THRESHOLD = 0.3

# -- Training Data --

KEYWORD_TO_SUBCATEGORIES = [
    # Suspension
    {"keywords": ["shock", "strut", "loaded strut", "quick strut"], "subcategories": ["Suspension > Shocks and Struts"]},
    {"keywords": ["control arm", "ball joint", "sway bar", "stabilizer bar", "stabilizer link", "tie rod", "track bar"], "subcategories": ["Suspension > Suspension, Springs and Related Components"]},
    {"keywords": ["coil spring", "leaf spring", "spring assembly"], "subcategories": ["Suspension > Suspension, Springs and Related Components"]},
    {"keywords": ["alignment", "camber", "caster", "toe"], "subcategories": ["Suspension > Alignment"]},
    # Brake
    {"keywords": ["brake pad", "brake shoe"], "subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    {"keywords": ["brake rotor", "brake drum", "rotor"], "subcategories": ["Brake > Drums and Rotors"]},
    {"keywords": ["brake caliper", "caliper bracket", "brake hose"], "subcategories": ["Brake > Brake Hydraulics"]},
    {"keywords": ["brake fluid", "brake bleed", "brake flush"], "subcategories": ["Brake > Brake System Service", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["abs system", "abs sensor", "abs module", "abs pump", "abs light", "anti-lock"], "subcategories": ["Brake > ABS Components"]},
    # Engine
    {"keywords": ["valve cover", "head gasket", "intake manifold gasket", "exhaust manifold gasket", "oil pan gasket"], "subcategories": ["Engine > Gaskets and Sealing Systems"]},
    {"keywords": ["timing chain", "timing belt", "camshaft", "crankshaft", "rocker arm", "lifter", "engine valve", "intake valve", "exhaust valve"], "subcategories": ["Engine > Valve Train Components"]},
    {"keywords": ["piston", "engine block", "connecting rod", "main bearing", "cylinder sleeve", "cylinder liner"], "subcategories": ["Engine > Cylinder Block Components"]},
    # Belts and Cooling
    {"keywords": ["serpentine", "drive belt", "belt tensioner", "idler pulley"], "subcategories": ["Belts and Cooling > Accessory Drive Belt System Components"]},
    {"keywords": ["radiator", "radiator hose"], "subcategories": ["Belts and Cooling > Radiators, Coolers and Related Components"]},
    {"keywords": ["coolant", "antifreeze"], "subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["water pump"], "subcategories": ["Belts and Cooling > Water Pump and Related Components"]},
    {"keywords": ["thermostat"], "subcategories": ["Belts and Cooling > Thermostat and Housing"]},
    {"keywords": ["cooling fan", "fan clutch", "fan motor", "radiator fan"], "subcategories": ["Belts and Cooling > Cooling Fan, Clutch and Motor"]},
    {"keywords": ["coolant hose", "heater hose"], "subcategories": ["Belts and Cooling > Hoses and Pipes"]},
    # Air and Fuel Delivery
    {"keywords": ["fuel pump", "fuel pump assembly", "fuel sender"], "subcategories": ["Air and Fuel Delivery > Fuel Pumps and Related Components"]},
    {"keywords": ["fuel injector", "fuel rail", "throttle body", "throttle"], "subcategories": ["Air and Fuel Delivery > Fuel Injection System and Related Components"]},
    {"keywords": ["air filter", "intake filter"], "subcategories": ["Air and Fuel Delivery > Filters"]},
    {"keywords": ["fuel tank", "fuel filler"], "subcategories": ["Air and Fuel Delivery > Fuel Storage"]},
    {"keywords": ["turbo", "turbocharger", "supercharger", "intercooler", "wastegate", "boost"], "subcategories": ["Air and Fuel Delivery > Turbocharger, Supercharger and Ram Air"]},
    # Electrical
    {"keywords": ["battery", "battery cable", "battery terminal"], "subcategories": ["Electrical, Charging and Starting > Battery and Related Components"]},
    {"keywords": ["alternator", "generator"], "subcategories": ["Electrical, Charging and Starting > Alternator / Generator and Related Components"]},
    {"keywords": ["starter", "starter motor"], "subcategories": ["Electrical, Charging and Starting > Starter and Related Components"]},
    # Emission Control
    {"keywords": ["catalytic converter", "cat converter"], "subcategories": ["Exhaust > Catalytic Converter"]},
    {"keywords": ["oxygen sensor", "o2 sensor"], "subcategories": ["Emission Control > Sensors"]},
    {"keywords": ["egr", "egr valve"], "subcategories": ["Emission Control > EGR and Related Components"]},
    {"keywords": ["evap", "purge valve", "charcoal canister", "vapor canister"], "subcategories": ["Emission Control > Evaporative System"]},
    {"keywords": ["pcv", "pcv valve"], "subcategories": ["Emission Control > Crankcase Ventilation System"]},
    # Exhaust
    {"keywords": ["muffler"], "subcategories": ["Exhaust > Mufflers"]},
    {"keywords": ["exhaust pipe", "exhaust manifold", "flex pipe", "resonator"], "subcategories": ["Exhaust > Exhaust and Tail Pipes", "Exhaust > Exhaust Manifolds"]},
    # Ignition
    {"keywords": ["spark plug"], "subcategories": ["Ignition > Secondary Ignition"]},
    {"keywords": ["ignition coil", "coil pack", "coil on plug"], "subcategories": ["Ignition > Ignition Coil"]},
    {"keywords": ["plug wire", "ignition wire"], "subcategories": ["Ignition > Ignition Wire and Related Components"]},
    {"keywords": ["distributor"], "subcategories": ["Ignition > Distributor and Magneto"]},
    # Steering
    {"keywords": ["power steering pump", "power steering hose", "power steering fluid", "power steering flush", "power steering"], "subcategories": ["Steering > Power Steering Hoses, Pumps, and Related Components", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["rack and pinion", "steering gear", "steering rack", "inner tie rod", "outer tie rod"], "subcategories": ["Steering > Steering, Gear and Related Components"]},
    {"keywords": ["steering column"], "subcategories": ["Steering > Steering Column"]},
    # Transmission
    {"keywords": ["transmission fluid", "trans fluid", "atf", "transmission flush"], "subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["clutch", "flywheel", "clutch disc", "pressure plate", "throwout bearing"], "subcategories": ["Transmission > Clutch Plates", "Transmission > Flywheel"]},
    {"keywords": ["torque converter"], "subcategories": ["Transmission > Automatic Transmission Components"]},
    # Driveline
    {"keywords": ["cv axle", "cv joint", "cv boot"], "subcategories": ["Driveline and Axles > CV Components", "Driveline and Axles > CV Boots and Related Components"]},
    {"keywords": ["drive shaft", "u-joint", "universal joint"], "subcategories": ["Driveline and Axles > Drive Shaft"]},
    {"keywords": ["differential", "diff fluid", "diff service"], "subcategories": ["Driveline and Axles > Differential", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["wheel bearing", "hub bearing", "hub assembly"], "subcategories": ["Driveline and Axles > Wheel Bearings, Seals, and Related Components", "Driveline and Axles > Hubs and Related Components"]},
    # Transfer Case
    {"keywords": ["transfer case"], "subcategories": ["Transfer Case > Transfer Case Components", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    # HVAC
    {"keywords": ["a/c compressor", "ac compressor"], "subcategories": ["HVAC > A/C Clutch and Compressor", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["a/c condenser", "ac condenser", "evaporator"], "subcategories": ["HVAC > A/C Condenser and Evaporator", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["refrigerant", "freon", "a/c recharge", "ac recharge", "a/c service", "ac service"], "subcategories": ["HVAC > A/C Clutch and Compressor", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    {"keywords": ["blower motor", "blower resistor"], "subcategories": ["HVAC > Motors, Core, Case and Related Components"]},
    {"keywords": ["cabin filter", "cabin air filter"], "subcategories": ["HVAC > Filters"]},
    {"keywords": ["heater core"], "subcategories": ["HVAC > Air Heater Components"]},
    # Tire and Wheel
    {"keywords": ["tire", "tyre"], "subcategories": ["Tire and Wheel > Tire"]},
    {"keywords": ["wheel", "rim", "lug nut"], "subcategories": ["Tire and Wheel > Wheel"]},
    {"keywords": ["tpms", "tire pressure sensor"], "subcategories": ["Tire and Wheel > Sensors"]},
    # Wiper and Washer
    {"keywords": ["wiper", "wiper blade"], "subcategories": ["Wiper and Washer > Wiper Arm, Blade, and Related Components"]},
    {"keywords": ["washer fluid", "washer pump", "windshield washer"], "subcategories": ["Wiper and Washer > Windshield Washer and Related Components"]},
    # Oil and Chemicals
    {"keywords": ["oil change", "oil service", "motor oil", "engine oil", "synthetic oil", "oil filter"], "subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)", "Engine > Filters"]},
    {"keywords": ["brake cleaner", "dielectric grease", "anti-seize", "thread locker", "penetrating oil"], "subcategories": ["Oil, Fluids and Chemicals > Cleaners and Removers", "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
]

VARIANT_TEMPLATES = [
    "Replace {kw}",
    "R&R {kw}",
    "{kw} replacement",
    "{kw} replacement needed",
    "Install new {kw}",
    "{kw} is bad",
    "{kw} failing",
    "{kw} worn out",
    "{kw} leaking",
    "{kw} noise",
    "needs {kw}",
    "new {kw}",
    "{kw} - replace",
    "Customer states {kw} issue",
    "Check and replace {kw}",
    "Recommend {kw} replacement",
]

SYMPTOM_VARIANTS = [
    # Brakes — pads/shoes
    {"symptoms": [
        "grinding noise when braking", "squealing when stopping", "brake pedal goes to floor",
        "pulsating brake pedal", "car pulls when braking", "scraping sound when stopping",
        "metal on metal when braking", "grinding from front when stopping",
        "high pitched squeal stopping", "brakes feel spongy",
    ], "subcategories": ["Brake > Disc Pads and Brake Shoes"]},
    # Brakes — rotors/drums
    {"symptoms": [
        "steering wheel shakes when braking", "vibration when stopping",
        "pulsation in pedal at highway speed", "shimmy when braking",
        "warped rotors vibration", "shudder when applying brakes",
    ], "subcategories": ["Brake > Drums and Rotors"]},
    {"symptoms": ["abs light on", "anti-lock brake warning", "abs fault code", "traction control light and abs light"], "subcategories": ["Brake > ABS Components"]},
    # Suspension — shocks/struts (expert-expanded)
    {"symptoms": [
        "clunking over bumps", "vehicle bounces excessively", "nose dives when stopping",
        "rough ride", "bouncy ride", "car keeps bouncing after bumps",
        "thump over speed bumps", "leaking strut", "front end bottoming out",
        "rear sags under load", "knocking over potholes", "harsh ride quality",
        "suspension feels loose", "banging noise hitting bumps",
        "front end clunk", "noise over railroad tracks",
        "bang over bumps", "thud from front end", "knock over rough road",
        "chassis slap over bumps", "double thump speed bumps",
        "front end bangs dips", "rear clunk over bumps", "trunk area thump bumps",
        "rattle over small bumps", "suspension noise over potholes",
        "clunk from front suspension", "rear suspension thump",
        "noise over expansion joints", "knocks on washboard road",
        "clunk over driveway apron", "noise from strut area over bumps",
    ], "subcategories": ["Suspension > Shocks and Struts"]},
    # Suspension — strut mounts (expert-expanded)
    {"symptoms": [
        "upper mount plate", "strut mount noise", "top mount bearing worn",
        "upper strut mount", "mount plate clunking", "popping when turning steering wheel stationary",
        "creaking from top of strut tower", "strut plate bearing",
        "upper strut plate", "mount plate bearing", "strut tower mount",
        "top mount worn", "upper spring seat", "strut top hat",
        "mount plate popping", "bearing plate noise turning",
        "strut mount bearing noisy", "upper mount clunk", "top hat binding",
        "spring seat cracked", "mount bearing seized",
        "pop while steering from strut mount", "front mount plate noise",
        "upper mount separated",
    ], "subcategories": ["Suspension > Shocks and Struts"]},
    # Suspension — alignment (expert-expanded)
    {"symptoms": [
        "vehicle pulls to one side", "uneven tire wear", "steering wheel off center",
        "vehicle pulls right", "vehicle pulls left", "car drifts right",
        "car drifts left", "steering wheel crooked", "feathering on tires",
        "cupping on tires", "car doesn't track straight", "pulls to the right on highway",
        "pulls to right", "drifts right on flat road", "steering pulls",
        "needs alignment pulls", "alignment off pulls right", "tracks right",
        "wanders on highway", "won't hold straight line",
        "off center steering wheel", "lane drift", "darting on highway",
    ], "subcategories": ["Suspension > Alignment"]},
    # Suspension — control arms/bushings (expert-expanded)
    {"symptoms": [
        "clunk when turning", "loose steering feel",
        "control arm bushing worn", "worn control arm bushings",
        "sway bar link rattle", "stabilizer bar noise",
        "ball joint play", "loose front end", "front end clunk turning",
        "bushing cracked", "bushing deteriorated", "lower control arm play",
        "LCA bushing torn", "compliance bushing shot",
        "control arm rear bushing split", "bushing separated",
        "arm bushing excessive movement", "thrust arm bushing cracked",
        "front end looseness low speed", "clunk turning and braking",
        "hydro bushing leaking", "rear compliance bushing torn",
        "FCA bushing worn", "RCA bushing worn",
    ], "subcategories": ["Suspension > Suspension, Springs and Related Components"]},
    # Engine — gaskets
    {"symptoms": ["oil leak from top of engine", "burning oil smell", "oil seeping from valve cover", "wet around intake"], "subcategories": ["Engine > Gaskets and Sealing Systems"]},
    # Engine — timing/valve train (expert-expanded)
    {"symptoms": [
        "rattling on cold start", "timing noise", "engine ticking",
        "timing chain replacement needed", "timing chain stretched",
        "timing chain rattle", "chain slap on startup", "variable valve timing code",
        "cam phaser noise", "vvt solenoid", "engine timing off",
        "chain stretched", "timing chain code", "P0016 cam crank correlation",
        "P0017 cam position", "timing set replacement", "chain guides worn",
        "cold start chain rattle", "chain noise on startup", "chain slack",
        "cam crank out of sync", "correlation fault", "chain tensioner weak",
        "guides broken", "phaser correlation fault",
        "cam timing over retarded", "cam timing over advanced",
        "engine rattles first start", "chain kit recommended",
        "front cover timing repair", "timing chain needed",
    ], "subcategories": ["Engine > Valve Train Components"]},
    # Cooling
    {"symptoms": ["overheating", "temperature gauge high", "coolant on ground", "coolant leak", "steam from hood"], "subcategories": ["Belts and Cooling > Radiators, Coolers and Related Components"]},
    {"symptoms": [
        "squealing belt noise", "belt chirping on startup",
        "serpentine belt cracking", "belt squeal", "chirping from front of engine",
        "belt slipping when wet", "cracked belt", "belt glazed",
    ], "subcategories": ["Belts and Cooling > Accessory Drive Belt System Components"]},
    # Electrical
    {"symptoms": ["car won't start", "clicking when turning key", "no crank", "slow crank", "starter grinds"], "subcategories": ["Electrical, Charging and Starting > Starter and Related Components"]},
    {"symptoms": ["battery light on", "dim headlights", "electrical issues", "charging system warning"], "subcategories": ["Electrical, Charging and Starting > Alternator / Generator and Related Components"]},
    {"symptoms": ["dead battery", "battery keeps dying", "battery won't hold charge", "needs jump every morning"], "subcategories": ["Electrical, Charging and Starting > Battery and Related Components"]},
    # Ignition — CEL/misfire (expert-expanded)
    {"symptoms": [
        "CEL on", "check engine light", "misfire", "rough idle", "rough running",
        "CEL on rough running", "engine misfiring", "stumbles at idle",
        "rough running check engine light", "cylinder misfire code",
        "P0300 random misfire", "P0301 misfire cylinder 1", "hesitation under acceleration",
        "tune up needed", "tune-up", "engine tune up",
        "CEL rough running", "CEL with misfire", "check engine rough idle",
        "MIL on running rough", "SES light rough idle", "CEL P0300",
        "CEL multiple misfire", "engine light on stumbling",
        "CEL on with shaking", "service engine soon rough",
        "misfiring under load", "flashing CEL misfire",
        "sputtering with check engine light", "shaking at idle CEL",
        "rough idle with codes", "stumble and check engine light",
        "P0301", "P0302", "P0303", "P0304",
        "random misfire", "multiple cylinder misfire",
        "bucks and jerks with CEL",
    ], "subcategories": ["Ignition > Secondary Ignition"]},
    # Ignition — spark plugs maintenance (expert-expanded)
    {"symptoms": [
        "spark plugs replacement due", "plugs due for service", "100k spark plug service",
        "plugs due", "spark plug service", "60k plug change",
        "100k plug service", "iridium plugs due", "platinum plugs worn",
        "plug gap too wide", "mileage-based plug replacement", "maintenance plugs",
        "tune up due", "due for plugs per mileage", "recommend spark plugs",
        "plugs worn", "electrodes worn", "plugs original",
        "overdue for spark plugs", "plug service due by maintenance schedule",
        "replace plugs as maintenance",
    ], "subcategories": ["Ignition > Secondary Ignition"]},
    # HVAC — A/C (expert-expanded)
    {"symptoms": [
        "a/c blowing warm", "no cold air", "a/c not working",
        "AC recharge needed blowing warm", "ac not cold enough",
        "ac recharge needed", "a/c weak", "a/c warm on one side",
        "refrigerant low", "ac blows hot", "air conditioning not cooling",
        "AC recharge", "A/C recharge needed", "freon low",
        "134a recharge", "1234yf recharge", "ac charge and check",
        "air not blowing cold", "AC performance check",
        "AC system evacuation and recharge", "AC evac and recharge",
        "low refrigerant", "static pressure low", "recharge and leak test",
        "recover evacuate recharge", "recharge with dye",
        "not cooling at idle", "vent temp high", "AC weak",
        "warm air from vents",
    ], "subcategories": ["HVAC > A/C Clutch and Compressor"]},
    {"symptoms": ["no heat", "heater not working", "heat blows cold", "lukewarm heat", "heater weak"], "subcategories": ["HVAC > Air Heater Components"]},
    # HVAC — blower motor (expert-expanded)
    {"symptoms": [
        "blower motor not working", "no air from vents", "fan not blowing",
        "blower only works on high", "hvac fan stopped", "no air flow from dash",
        "blower resistor", "fan speed not changing", "only high speed works on fan",
        "HVAC fan dead", "cabin fan stopped", "no air movement dash",
        "blower quit working", "heat fan not blowing", "fan motor seized",
        "blower motor intermittent", "no fan from vents", "blower inop",
        "interior fan not working", "no blower on any speed",
        "blower resistor concern", "fan cuts in and out",
        "no airflow from vents", "HVAC blower intermittent",
    ], "subcategories": ["HVAC > Motors, Core, Case and Related Components"]},
    # Exhaust — muffler
    {"symptoms": ["loud exhaust", "exhaust leak", "rumbling underneath", "exhaust drone", "holes in muffler"], "subcategories": ["Exhaust > Mufflers"]},
    # Exhaust — catalytic converter (expert-expanded)
    {"symptoms": [
        "rotten egg smell", "catalytic converter code", "sulfur smell from exhaust",
        "P0420 catalyst efficiency", "P0430 catalyst", "cat efficiency below threshold",
        "rotten egg odor", "exhaust smells like eggs", "failing catalytic converter",
        "egg smell from exhaust", "sulfur odor driving", "rotten smell tailpipe",
        "hydrogen sulfide smell", "cat going bad smell", "stinky exhaust",
        "exhaust smells terrible", "rotten smell after hard acceleration",
        "bad catalyst smell", "converter smell", "strong sulfur exhaust odor",
        "rotten egg smell under load", "sulfur smell after highway drive",
    ], "subcategories": ["Exhaust > Catalytic Converter"]},
    # Steering — power steering (expert-expanded)
    {"symptoms": [
        "whining when turning", "hard to steer", "power steering whine",
        "whining noise turning wheel", "groan when turning steering wheel",
        "steering stiff", "fluid leaking from steering rack",
        "whine at low speed turns", "moaning sound turning",
        "power steering fluid low", "ps pump noise",
        "whine at full lock", "pump noise turning", "moan turning wheel",
        "steering whine low speed", "PS pump whine", "groan turning at idle",
        "howling when steering", "fluid cavitation noise turning",
        "power steering noisy", "steering pump moans",
        "whine when parking", "groan backing out while turning",
        "steering noise cold", "aerated PS fluid", "low PS fluid noise",
    ], "subcategories": ["Steering > Power Steering Hoses, Pumps, and Related Components"]},
    # Driveline — CV
    {"symptoms": ["clicking when turning", "cv clicking", "clicking noise in turns", "cv axle popping"], "subcategories": ["Driveline and Axles > CV Components"]},
    # Driveline — wheel bearings (expert-expanded)
    {"symptoms": [
        "humming from wheel area", "wheel noise at speed", "growling noise from front",
        "humming gets louder with speed", "bearing noise",
        "roaring sound from wheel", "drone changes with steering input",
        "wheel hum goes away turning one direction", "rumble from rear axle area",
        "hub bearing noise", "grinding from front wheel",
        "hum from front wheel", "droning noise road speed", "wheel rumble highway",
        "noise changes turning left right", "growl from hub area",
        "bearing drone 40mph plus", "hum louder faster you go",
        "helicopter noise road speed", "wheel bearing growl",
        "drone from left front", "drone from right front",
        "noise unloads on turns", "hub noise",
        "rumble from corner of vehicle",
    ], "subcategories": ["Driveline and Axles > Wheel Bearings, Seals, and Related Components"]},
    # Body — mirrors
    {"symptoms": [
        "mirror assembly", "mirror replacement", "side mirror broken",
        "mirror glass cracked", "power mirror not working", "mirror folding mechanism",
        "driver side mirror", "passenger mirror", "rear view mirror fell off",
        "mirror turn signal", "blind spot mirror",
    ], "subcategories": ["Body > Mirrors"]},
    # Tire — TPMS (expert-expanded)
    {"symptoms": [
        "TPMS light on", "tire pressure light", "tire pressure warning",
        "tpms sensor", "tire monitor light", "low tire pressure warning",
        "tpms fault", "tire pressure monitoring light",
        "tire light on dash", "TPMS warning", "tire monitor fault",
        "low tire light stays on after fill", "TPMS sensor battery dead",
        "tire pressure indicator", "TPMS relearn needed",
        "tire pressure light flashing", "sensor not detected",
        "tire pressure monitor inop", "tpms service",
        "no sensor signal", "relearn tire sensors", "sensor battery low",
    ], "subcategories": ["Tire and Wheel > Sensors"]},
    # Fluids — ATF (expert-expanded)
    {"symptoms": [
        "ATF flush", "ATF change", "transmission fluid service",
        "trans fluid flush", "automatic transmission fluid",
        "transmission fluid dark", "burnt trans fluid",
        "ATF exchange", "ATF drain and fill", "automatic transmission fluid service",
        "transmission service ATF", "ATF dark needs flush",
        "ATF maintenance", "trans fluid exchange",
        "transmission fluid service due", "CVT fluid service",
        "DSG service", "pan drop and filter", "trans service by mileage",
        "burnt ATF", "fluid discolored", "fluid exchange recommended",
        "service transmission fluid", "trans fluid change",
    ], "subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
    # Fluids — oil change with marketing prefix
    {"symptoms": [
        "Mailer oil change", "oil change special", "coupon oil change",
        "promo oil change", "mail-in oil change", "oil change deal",
        "customer came in for oil change", "just here for an oil change",
        "routine oil change", "express oil change", "quick lube oil change",
    ], "subcategories": ["Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)"]},
]

SYMPTOM_TEMPLATES = [
    "{symptom}",
    "Customer states {symptom}",
    "C/S {symptom}",
    "Complaint: {symptom}",
    "Concern: {symptom}",
    "Vehicle has {symptom}",
    "Car {symptom}",
    "Inspect for {symptom}",
    "Diag for {symptom}",
    "Check {symptom}",
    # Expert-added: DMS/shop shorthand prefixes
    "Cust states {symptom}",
    "States {symptom}",
    "Complains of {symptom}",
    "Chk and advise {symptom}",
    "Check and advise {symptom}",
    "Inspect and advise {symptom}",
    "Found {symptom}",
    "Verified {symptom}",
    "Confirmed {symptom}",
    "Recommend {symptom}",
    "Rec {symptom}",
    "LF {symptom}",
    "RF {symptom}",
    "LR {symptom}",
    "RR {symptom}",
    "Front {symptom}",
    "Rear {symptom}",
]

COMMON_TYPOS = {
    "brake": ["brak", "brke", "barke"],
    "rotor": ["roter", "rotot"],
    "strut": ["strutt", "strrt"],
    "alternator": ["alternater", "alterntor"],
    "thermostat": ["thermistat", "thermstat", "thermastat"],
    "serpentine": ["serpantine", "serpentene"],
    "catalytic": ["catalitic", "catylitic", "catalyic"],
    "differential": ["differntial", "diferential"],
    "compressor": ["compresser", "compresser"],
    "radiator": ["raditor", "raidator"],
    "transmission": ["transmision", "trnasmission"],
    "steering": ["steerin", "stearing"],
    "exhaust": ["exaust", "exhuast"],
    "alignment": ["alignemnt", "aligment"],
    "battery": ["baterry", "batter"],
    "coolant": ["coolent", "coollant"],
    "muffler": ["mufler", "mufflar"],
}

HARD_NEGATIVES = [
    # --- Original hard negatives ---
    {"query": "replace radiator", "positive": "Belts and Cooling > Radiators, Coolers and Related Components", "negative": "Transmission > Radiators, Coolers and Related Components"},
    {"query": "r&r radiator", "positive": "Belts and Cooling > Radiators, Coolers and Related Components", "negative": "Transmission > Radiators, Coolers and Related Components"},
    {"query": "radiator leaking", "positive": "Belts and Cooling > Radiators, Coolers and Related Components", "negative": "Transmission > Radiators, Coolers and Related Components"},
    {"query": "radiator hose replacement", "positive": "Belts and Cooling > Radiators, Coolers and Related Components", "negative": "Transmission > Radiators, Coolers and Related Components"},
    {"query": "transmission cooler", "positive": "Transmission > Radiators, Coolers and Related Components", "negative": "Belts and Cooling > Radiators, Coolers and Related Components"},
    {"query": "trans cooler line leaking", "positive": "Transmission > Radiators, Coolers and Related Components", "negative": "Belts and Cooling > Radiators, Coolers and Related Components"},
    {"query": "timing chain replacement", "positive": "Engine > Valve Train Components", "negative": "Driveline and Axles > Drive Shaft"},
    {"query": "timing belt worn", "positive": "Engine > Valve Train Components", "negative": "Belts and Cooling > Accessory Drive Belt System Components"},
    {"query": "spark plug replacement", "positive": "Ignition > Secondary Ignition", "negative": "Ignition > Ignition Coil"},
    {"query": "new spark plugs", "positive": "Ignition > Secondary Ignition", "negative": "Ignition > Ignition Coil"},
    {"query": "CEL on rough running", "positive": "Ignition > Secondary Ignition", "negative": "Electrical, Charging and Starting > Battery and Related Components"},
    {"query": "check engine light misfire", "positive": "Ignition > Secondary Ignition", "negative": "Emission Control > Sensors"},
    {"query": "ATF flush", "positive": "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)", "negative": "Transmission > Automatic Transmission Components"},
    {"query": "transmission fluid change", "positive": "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)", "negative": "Transmission > Automatic Transmission Components"},
    {"query": "Mailer oil change", "positive": "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)", "negative": "Engine > Gaskets and Sealing Systems"},
    {"query": "oil change special", "positive": "Oil, Fluids and Chemicals > Functional Fluid, Lubricant, Grease (including Additives)", "negative": "Engine > Gaskets and Sealing Systems"},
    {"query": "clunking over bumps", "positive": "Suspension > Shocks and Struts", "negative": "Brake > Drums and Rotors"},
    {"query": "clunking over bumps", "positive": "Suspension > Shocks and Struts", "negative": "Driveline and Axles > CV Boots and Related Components"},
    {"query": "noise over speed bumps", "positive": "Suspension > Shocks and Struts", "negative": "Transmission > Clutch Plates"},
    {"query": "rotten egg smell", "positive": "Exhaust > Catalytic Converter", "negative": "Engine > Valve Train Components"},
    {"query": "rotten egg smell", "positive": "Exhaust > Catalytic Converter", "negative": "Emission Control > Evaporative System"},
    {"query": "sulfur smell exhaust", "positive": "Exhaust > Catalytic Converter", "negative": "Engine > Gaskets and Sealing Systems"},
    {"query": "whining when turning", "positive": "Steering > Power Steering Hoses, Pumps, and Related Components", "negative": "Emission Control > Crankcase Ventilation System"},
    {"query": "whining when turning", "positive": "Steering > Power Steering Hoses, Pumps, and Related Components", "negative": "Transmission > Flywheel"},
    {"query": "whine turning wheel", "positive": "Steering > Power Steering Hoses, Pumps, and Related Components", "negative": "Emission Control > Sensors"},
    {"query": "humming from wheel area", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Tire and Wheel > Wheel"},
    {"query": "humming from wheel area", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Transmission > Flywheel"},
    {"query": "humming gets louder with speed", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Tire and Wheel > Sensors"},
    {"query": "vehicle pulls right", "positive": "Suspension > Alignment", "negative": "Transmission > Flywheel"},
    {"query": "vehicle pulls right", "positive": "Suspension > Alignment", "negative": "Driveline and Axles > Drive Shaft"},
    {"query": "car drifts right", "positive": "Suspension > Alignment", "negative": "Steering > Steering Column"},
    {"query": "mirror assembly", "positive": "Body > Mirrors", "negative": "Wiper and Washer > Windshield Washer and Related Components"},
    {"query": "mirror assembly", "positive": "Body > Mirrors", "negative": "Wiper and Washer > Wiper Arm, Blade, and Related Components"},
    {"query": "side mirror replacement", "positive": "Body > Mirrors", "negative": "Suspension > Alignment"},
    {"query": "blower motor not working", "positive": "HVAC > Motors, Core, Case and Related Components", "negative": "Exhaust > Catalytic Converter"},
    {"query": "blower motor not working", "positive": "HVAC > Motors, Core, Case and Related Components", "negative": "Emission Control > Crankcase Ventilation System"},
    {"query": "no air from vents", "positive": "HVAC > Motors, Core, Case and Related Components", "negative": "Air and Fuel Delivery > Turbocharger, Supercharger and Ram Air"},
    {"query": "upper mount plate", "positive": "Suspension > Shocks and Struts", "negative": "Transmission > Clutch Plates"},
    {"query": "upper mount plate", "positive": "Suspension > Shocks and Struts", "negative": "Exhaust > Exhaust Manifolds"},
    {"query": "strut mount plate", "positive": "Suspension > Shocks and Struts", "negative": "Suspension > Alignment"},
    {"query": "control arm bushing worn", "positive": "Suspension > Suspension, Springs and Related Components", "negative": "Wiper and Washer > Wiper Arm, Blade, and Related Components"},
    {"query": "control arm bushing worn", "positive": "Suspension > Suspension, Springs and Related Components", "negative": "Emission Control > Crankcase Ventilation System"},
    {"query": "AC recharge needed blowing warm", "positive": "HVAC > A/C Clutch and Compressor", "negative": "Emission Control > Crankcase Ventilation System"},
    {"query": "ac not blowing cold", "positive": "HVAC > A/C Clutch and Compressor", "negative": "Transmission > Radiators, Coolers and Related Components"},
    {"query": "TPMS light on", "positive": "Tire and Wheel > Sensors", "negative": "Ignition > Distributor and Magneto"},
    {"query": "TPMS light on", "positive": "Tire and Wheel > Sensors", "negative": "Driveline and Axles > CV Boots and Related Components"},
    {"query": "tire pressure light on", "positive": "Tire and Wheel > Sensors", "negative": "Driveline and Axles > Wheel Bearings, Seals, and Related Components"},
    # --- Expert-added hard negatives (with corrected PCAdb names) ---
    {"query": "CEL rough running", "positive": "Ignition > Secondary Ignition", "negative": "Emission Control > Sensors"},
    {"query": "CEL rough running", "positive": "Ignition > Secondary Ignition", "negative": "Engine > Valve Train Components"},
    {"query": "CEL rough running", "positive": "Ignition > Secondary Ignition", "negative": "Air and Fuel Delivery > Fuel Pumps and Related Components"},
    {"query": "CEL rough running", "positive": "Ignition > Secondary Ignition", "negative": "Air and Fuel Delivery > Filters"},
    {"query": "Spark plugs replacement due", "positive": "Ignition > Secondary Ignition", "negative": "Ignition > Ignition Coil"},
    {"query": "Spark plugs replacement due", "positive": "Ignition > Secondary Ignition", "negative": "Engine > Valve Train Components"},
    {"query": "Spark plugs replacement due", "positive": "Ignition > Secondary Ignition", "negative": "Air and Fuel Delivery > Fuel Injection System and Related Components"},
    {"query": "AC recharge needed", "positive": "HVAC > A/C Clutch and Compressor", "negative": "HVAC > A/C Condenser and Evaporator"},
    {"query": "AC recharge needed", "positive": "HVAC > A/C Clutch and Compressor", "negative": "HVAC > Motors, Core, Case and Related Components"},
    {"query": "blower not working", "positive": "HVAC > Motors, Core, Case and Related Components", "negative": "HVAC > A/C Clutch and Compressor"},
    {"query": "blower not working", "positive": "HVAC > Motors, Core, Case and Related Components", "negative": "Belts and Cooling > Cooling Fan, Clutch and Motor"},
    {"query": "TPMS light", "positive": "Tire and Wheel > Sensors", "negative": "Tire and Wheel > Tire"},
    {"query": "timing chain needed", "positive": "Engine > Valve Train Components", "negative": "Belts and Cooling > Accessory Drive Belt System Components"},
    {"query": "control arm bushing", "positive": "Suspension > Suspension, Springs and Related Components", "negative": "Suspension > Shocks and Struts"},
    {"query": "control arm bushing", "positive": "Suspension > Suspension, Springs and Related Components", "negative": "Steering > Steering, Gear and Related Components"},
    {"query": "clunking over bumps", "positive": "Suspension > Shocks and Struts", "negative": "Suspension > Suspension, Springs and Related Components"},
    {"query": "clunking over bumps", "positive": "Suspension > Shocks and Struts", "negative": "Steering > Steering, Gear and Related Components"},
    {"query": "clunking over bumps", "positive": "Suspension > Shocks and Struts", "negative": "Brake > Brake Hydraulics"},
    {"query": "vehicle pulls right", "positive": "Suspension > Alignment", "negative": "Tire and Wheel > Tire"},
    {"query": "vehicle pulls right", "positive": "Suspension > Alignment", "negative": "Brake > Brake Hydraulics"},
    {"query": "vehicle pulls right", "positive": "Suspension > Alignment", "negative": "Suspension > Suspension, Springs and Related Components"},
    {"query": "whining when turning", "positive": "Steering > Power Steering Hoses, Pumps, and Related Components", "negative": "Belts and Cooling > Accessory Drive Belt System Components"},
    {"query": "whining when turning", "positive": "Steering > Power Steering Hoses, Pumps, and Related Components", "negative": "Driveline and Axles > CV Components"},
    {"query": "humming from wheel area", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Tire and Wheel > Tire"},
    {"query": "humming from wheel area", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Brake > Brake Hydraulics"},
    {"query": "humming from wheel area", "positive": "Driveline and Axles > Wheel Bearings, Seals, and Related Components", "negative": "Driveline and Axles > Differential"},
]


# ============================================================================
# DATA GENERATION
# ============================================================================

def apply_typo(text: str) -> str | None:
    lower = text.lower()
    for correct, typos in COMMON_TYPOS.items():
        if correct in lower:
            typo = random.choice(typos)
            return text.replace(correct, typo).replace(correct.capitalize(), typo.capitalize())
    return None


def generate_training_data() -> tuple[list[dict], list[dict]]:
    """Generate training pairs and triplets. Returns (positive_pairs, triplets)."""
    qp = "search_query: "
    dp = "search_document: "
    positive_pairs = []
    triplets = []

    # Source 1: keyword → subcategory with variant templates
    for entry in KEYWORD_TO_SUBCATEGORIES:
        for kw in entry["keywords"]:
            for subcat in entry["subcategories"]:
                for tmpl in VARIANT_TEMPLATES:
                    query = tmpl.format(kw=kw)
                    positive_pairs.append({"query": f"{qp}{query}", "positive": f"{dp}{subcat}"})
                typo_text = apply_typo(f"Replace {kw}")
                if typo_text:
                    positive_pairs.append({"query": f"{qp}{typo_text}", "positive": f"{dp}{subcat}"})

    # Source 2: symptom-based with template expansion
    for entry in SYMPTOM_VARIANTS:
        for symptom in entry["symptoms"]:
            for subcat in entry["subcategories"]:
                for tmpl in SYMPTOM_TEMPLATES:
                    query = tmpl.format(symptom=symptom)
                    positive_pairs.append({"query": f"{qp}{query}", "positive": f"{dp}{subcat}"})

    # Source 3: hard negatives (triplets)
    for hn in HARD_NEGATIVES:
        triplets.append({
            "query": f"{qp}{hn['query']}",
            "positive": f"{dp}{hn['positive']}",
            "negative": f"{dp}{hn['negative']}",
        })

    random.shuffle(positive_pairs)
    random.shuffle(triplets)
    return positive_pairs, triplets


# ============================================================================
# EVALUATION — DO NOT MODIFY
# ============================================================================

def strip_prefix(text: str) -> str:
    for prefix in ("search_query: ", "search_document: "):
        if text.startswith(prefix):
            return text[len(prefix):]
    return text


def evaluate_model(model: SentenceTransformer, corpus_names: list[str]) -> dict:
    """Run benchmark evaluation. Returns metrics dict."""
    k_values = [1, 3, 5, 10]

    corpus_texts = [f"search_document: {name}" for name in corpus_names]
    corpus_embeddings = model.encode(corpus_texts, normalize_embeddings=True, show_progress_bar=False)

    query_texts = [f"search_query: {b['concern']}" for b in BENCHMARK]
    query_embeddings = model.encode(query_texts, normalize_embeddings=True, show_progress_bar=False)

    sim_matrix = np.array(query_embeddings) @ np.array(corpus_embeddings).T

    hits_at_k = {k: 0 for k in k_values}
    failures = []
    garbage_correct = 0
    garbage_total = 0

    for i, case in enumerate(BENCHMARK):
        expected = set(case["expected_subcategories"])
        scores = sim_matrix[i]
        ranked_indices = np.argsort(scores)[::-1]

        if len(expected) == 0:
            garbage_total += 1
            top_score = scores[ranked_indices[0]]
            if top_score < GARBAGE_THRESHOLD:
                garbage_correct += 1
            continue

        for k in k_values:
            top_k_names = {corpus_names[idx] for idx in ranked_indices[:k]}
            if expected & top_k_names:
                hits_at_k[k] += 1

        top_5_names = [corpus_names[idx] for idx in ranked_indices[:5]]
        if not (expected & set(top_5_names)):
            failures.append(case["concern"])

    n_real = sum(1 for c in BENCHMARK if len(c["expected_subcategories"]) > 0)
    accuracy = {k: hits_at_k[k] / n_real * 100 for k in k_values}

    return {
        "hit_at_1": accuracy[1],
        "hit_at_5": accuracy[5],
        "hit_at_10": accuracy[10],
        "garbage_correct": garbage_correct,
        "garbage_total": garbage_total,
        "failures": failures,
        "n_real": n_real,
    }


# ============================================================================
# TRAINING
# ============================================================================

def train():
    import torch
    print(f"=== Automotive Embedding Fine-Tuning ===")
    print(f"CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"Base model: {BASE_MODEL}")
    print(f"Epochs: {EPOCHS}, Batch size: {BATCH_SIZE}, LR: {LEARNING_RATE}")
    print(f"Matryoshka dims: {MATRYOSHKA_DIMS}")
    print()

    # Generate data
    print("Generating training data...")
    positive_pairs, triplets = generate_training_data()
    print(f"  Positive pairs: {len(positive_pairs)}")
    print(f"  Hard negative triplets: {len(triplets)}")

    # Build corpus for evaluation
    all_subcats = set()
    for p in positive_pairs:
        all_subcats.add(strip_prefix(p["positive"]))
    for t in triplets:
        all_subcats.add(strip_prefix(t["positive"]))
        all_subcats.add(strip_prefix(t["negative"]))
    corpus_names = sorted(all_subcats)
    print(f"  Corpus size: {len(corpus_names)} subcategories")

    # Load model
    print(f"\nLoading {BASE_MODEL}...")
    model = SentenceTransformer(BASE_MODEL, trust_remote_code=True)

    # Baseline evaluation
    print("\n--- Baseline (before training) ---")
    baseline = evaluate_model(model, corpus_names)
    print(f"  Hit@1={baseline['hit_at_1']:.1f}%  Hit@5={baseline['hit_at_5']:.1f}%  Hit@10={baseline['hit_at_10']:.1f}%")
    print(f"  Garbage: {baseline['garbage_correct']}/{baseline['garbage_total']}")
    if baseline["failures"]:
        print(f"  Failures ({len(baseline['failures'])}): {baseline['failures']}")

    # Build training objectives — MNRL (no Matryoshka) + hard negative TripletLoss
    pos_examples = [InputExample(texts=[p["query"], p["positive"]]) for p in positive_pairs]
    trip_examples = [InputExample(texts=[t["query"], t["positive"], t["negative"]]) for t in triplets]

    train_objectives = []

    if pos_examples:
        pos_loader = DataLoader(pos_examples, shuffle=True, batch_size=BATCH_SIZE)
        mnrl_loss = losses.MultipleNegativesRankingLoss(model)
        if MATRYOSHKA_DIMS:
            mnrl_loss = losses.MatryoshkaLoss(model, mnrl_loss, matryoshka_dims=MATRYOSHKA_DIMS)
        train_objectives.append((pos_loader, mnrl_loss))

    if trip_examples:
        trip_loader = DataLoader(trip_examples, shuffle=True, batch_size=BATCH_SIZE)
        triplet_loss = losses.TripletLoss(model)
        if MATRYOSHKA_DIMS:
            triplet_loss = losses.MatryoshkaLoss(model, triplet_loss, matryoshka_dims=MATRYOSHKA_DIMS)
        train_objectives.append((trip_loader, triplet_loss))

    total_steps = sum(len(loader) for loader, _ in train_objectives) * EPOCHS
    warmup_steps = int(total_steps * WARMUP_RATIO)

    # Train to a temp directory
    output_dir = tempfile.mkdtemp(prefix="automotive-embed-")
    print(f"\nTraining ({total_steps} steps, {warmup_steps} warmup)...")

    model.fit(
        train_objectives=train_objectives,
        epochs=EPOCHS,
        warmup_steps=warmup_steps,
        optimizer_params={"lr": LEARNING_RATE},
        output_path=output_dir,
        show_progress_bar=True,
    )

    # Evaluate the in-memory trained model directly
    # (Reloading from disk causes _IncompatibleKeys with nomic's custom architecture)
    print("\n--- Fine-tuned (after training) ---")
    result = evaluate_model(model, corpus_names)
    print(f"  Hit@1={result['hit_at_1']:.1f}%  Hit@5={result['hit_at_5']:.1f}%  Hit@10={result['hit_at_10']:.1f}%")
    print(f"  Garbage: {result['garbage_correct']}/{result['garbage_total']}")
    if result["failures"]:
        print(f"  Failures ({len(result['failures'])}): {result['failures']}")

    # Delta
    d1 = result["hit_at_1"] - baseline["hit_at_1"]
    d5 = result["hit_at_5"] - baseline["hit_at_5"]
    d10 = result["hit_at_10"] - baseline["hit_at_10"]
    print(f"\n--- Delta ---")
    print(f"  Hit@1: {'+' if d1 >= 0 else ''}{d1:.1f}%  Hit@5: {'+' if d5 >= 0 else ''}{d5:.1f}%  Hit@10: {'+' if d10 >= 0 else ''}{d10:.1f}%")

    # Parseable result line
    print(f"\nRESULT: hit@1={result['hit_at_1']:.1f} hit@5={result['hit_at_5']:.1f} hit@10={result['hit_at_10']:.1f} garbage={result['garbage_correct']}/{result['garbage_total']}")

    # Save model path
    print(f"\nModel saved to: {output_dir}")

    return result


if __name__ == "__main__":
    train()