Deploy CodeBERT training Space

.gitignore

@@ -0,0 +1,13 @@
+__pycache__/
+*.py[cod]
+.venv/
+venv/
+.env
+data/*.parquet
+data/*.csv
+models/*.joblib
+models/evaluation/
+.DS_Store
+*.egg-info/
+dist/
+build/

README.md

@@ -1,13 +1,46 @@
 ---
-title: Sql Error Classifier Train
-emoji: 📉
-colorFrom: gray
-colorTo: gray
+title: SQL Error Classifier Training
+emoji: 🧠
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 6.17.3
-python_version: '3.13'
-app_file: app.py
+sdk_version: 4.44.0
+app_file: train_space_app.py
 pinned: false
+license: mit
+hardware: t4-small
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SQL Error Classifier — CodeBERT Training Space
+
+Train `microsoft/codebert-base` as a **cross-encoder** for multi-label SQL error classification.
+
+## Setup
+
+1. **Hardware:** Settings → Hardware → **GPU t4-small** (recommended)
+2. **Secrets:** Settings → Secrets → add `HF_TOKEN` (Hugging Face write token) to push models to your account
+3. **Data:** Include `data/sql_errors_dev.parquet` in this Space repo, or upload parquet at runtime
+
+## Usage
+
+1. Choose bundled dataset or upload your own parquet
+2. Set epochs, batch size, max samples
+3. Click **Start Training**
+4. Optionally enable **Push to Hub** with model id `your-username/sql-codebert-classifier`
+
+## Dataset columns
+
+Required (aliases supported):
+
+| Column | Aliases |
+|--------|---------|
+| `question` | — |
+| `schema` | — |
+| `student_sql` | `query` |
+| `correct_sql` | `correct_query` |
+| `error_labels` | `label_name` |
+
+## Labels (9-class multi-label)
+
+`JOIN_ERROR`, `AGGREGATION_ERROR`, `FILTER_ERROR`, `WINDOW_FUNCTION_ERROR`,
+`SUBQUERY_ERROR`, `NULL_HANDLING_ERROR`, `PERFORMANCE_ERROR`, `LOGICAL_ERROR`, `SYNTAX_ERROR`

README_HF_SPACE.md

@@ -0,0 +1,17 @@
+---
+title: SQL Error Classifier
+emoji: 🗄️
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.44.0
+app_file: app.py
+pinned: false
+license: mit
+---
+
+# SQL Error Classifier
+
+Demo Space for the multi-tower SQL error classification model.
+
+Set `SPACE_MODEL_ID` in Space secrets to your model repo (e.g. `username/sql-error-classifier`).

README_TRAIN_SPACE.md

@@ -0,0 +1,46 @@
+---
+title: SQL Error Classifier Training
+emoji: 🧠
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 4.44.0
+app_file: train_space_app.py
+pinned: false
+license: mit
+hardware: t4-small
+---
+
+# SQL Error Classifier — CodeBERT Training Space
+
+Train `microsoft/codebert-base` as a **cross-encoder** for multi-label SQL error classification.
+
+## Setup
+
+1. **Hardware:** Settings → Hardware → **GPU t4-small** (recommended)
+2. **Secrets:** Settings → Secrets → add `HF_TOKEN` (Hugging Face write token) to push models to your account
+3. **Data:** Include `data/sql_errors_dev.parquet` in this Space repo, or upload parquet at runtime
+
+## Usage
+
+1. Choose bundled dataset or upload your own parquet
+2. Set epochs, batch size, max samples
+3. Click **Start Training**
+4. Optionally enable **Push to Hub** with model id `your-username/sql-codebert-classifier`
+
+## Dataset columns
+
+Required (aliases supported):
+
+| Column | Aliases |
+|--------|---------|
+| `question` | — |
+| `schema` | — |
+| `student_sql` | `query` |
+| `correct_sql` | `correct_query` |
+| `error_labels` | `label_name` |
+
+## Labels (9-class multi-label)
+
+`JOIN_ERROR`, `AGGREGATION_ERROR`, `FILTER_ERROR`, `WINDOW_FUNCTION_ERROR`,
+`SUBQUERY_ERROR`, `NULL_HANDLING_ERROR`, `PERFORMANCE_ERROR`, `LOGICAL_ERROR`, `SYNTAX_ERROR`

app.py

@@ -0,0 +1,156 @@
+"""
+Gradio app for Hugging Face Spaces.
+
+Deploy: create a Space with sdk=gradio and point app_file to this file.
+Set SPACE_MODEL_ID env var to your HF model repo (e.g. username/sql-error-classifier).
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import gradio as gr
+
+from src.huggingface import SQLLErrorClassifierHF
+
+MODEL_ID = os.getenv("SPACE_MODEL_ID", "models/hf_package")
+LOCAL_PACKAGE = Path(__file__).parent / "models" / "hf_package"
+
+EXAMPLE = {
+    "question": "What is the average score of students in each department?",
+    "schema": "students(id, name, score, department_id) | departments(id, name)",
+    "correct_query": (
+        "SELECT department_id, AVG(score) FROM students GROUP BY department_id"
+    ),
+    "student_query": (
+        "SELECT department_id, SUM(score) FROM students GROUP BY department_id"
+    ),
+    "error_message": "query executes but produces incorrect result set",
+}
+
+
+def _load_classifier() -> SQLLErrorClassifierHF:
+    if LOCAL_PACKAGE.exists() and (LOCAL_PACKAGE / "config.json").exists():
+        return SQLLErrorClassifierHF.from_pretrained(LOCAL_PACKAGE)
+    return SQLLErrorClassifierHF.from_pretrained(MODEL_ID)
+
+
+clf = _load_classifier()
+
+
+def classify(
+    question: str,
+    schema: str,
+    correct_query: str,
+    student_query: str,
+    error_message: str,
+) -> tuple[str, str, str]:
+    result = clf.predict(
+        question=question.strip(),
+        schema=schema.strip(),
+        correct_query=correct_query.strip(),
+        student_query=student_query.strip(),
+        error_message=error_message.strip() or None,
+    )
+
+    summary = (
+        f"**{result['label_name']}**  \n"
+        f"Confidence: **{result['confidence']:.1%}**"
+    )
+    top_k = "\n".join(
+        f"- {item['label_name']}: {item['confidence']:.1%}"
+        for item in result["top_k"]
+    )
+    sims = result.get("similarities") or result.get("pair_scores") or {}
+    diagnostics = "\n".join(
+        f"- **{k.replace('_', ' ').title()}**: {v:.3f}"
+        for k, v in sims.items()
+    )
+    if not diagnostics:
+        diagnostics = "_No diagnostic scores for this model type._"
+    return summary, top_k, diagnostics
+
+
+with gr.Blocks(title="SQL Error Classifier") as demo:
+    gr.Markdown(
+        """
+        # SQL Error Classifier
+        Classify **which mistake area** a student is struggling with, using:
+        **question**, **schema**, **correct query**, and the **student's query**.
+
+        Powered by a multi-tower MiniLM architecture on Hugging Face.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column():
+            question = gr.Textbox(
+                label="Question",
+                lines=2,
+                value=EXAMPLE["question"],
+            )
+            schema = gr.Textbox(
+                label="Schema",
+                lines=2,
+                value=EXAMPLE["schema"],
+            )
+            correct_query = gr.Textbox(
+                label="Correct Query",
+                lines=3,
+                value=EXAMPLE["correct_query"],
+            )
+            student_query = gr.Textbox(
+                label="Student Query",
+                lines=3,
+                value=EXAMPLE["student_query"],
+            )
+            error_message = gr.Textbox(
+                label="DB Error Message (optional)",
+                lines=2,
+                value=EXAMPLE["error_message"],
+            )
+            run_btn = gr.Button("Classify", variant="primary")
+
+        with gr.Column():
+            prediction = gr.Markdown(label="Prediction")
+            top_k = gr.Markdown(label="Top 3")
+            diagnostics = gr.Markdown(label="Semantic Diagnostics")
+
+    run_btn.click(
+        classify,
+        inputs=[question, schema, correct_query, student_query, error_message],
+        outputs=[prediction, top_k, diagnostics],
+    )
+
+    gr.Examples(
+        examples=[
+            [
+                EXAMPLE["question"],
+                EXAMPLE["schema"],
+                EXAMPLE["correct_query"],
+                EXAMPLE["student_query"],
+                EXAMPLE["error_message"],
+            ],
+            [
+                "Find students who have not provided an email address.",
+                "students(id, name, email, phone)",
+                "SELECT name FROM students WHERE email IS NULL",
+                "SELECT name FROM students WHERE email = NULL",
+                "use IS NULL or IS NOT NULL to test for null values",
+            ],
+            [
+                "List each student's name along with their department name.",
+                "students(id, name, department_id) | departments(id, name)",
+                "SELECT students.name, departments.name FROM students "
+                "INNER JOIN departments ON students.department_id = departments.id",
+                "SELECT students.name, departments.name FROM students JOIN departments",
+                "missing ON clause or invalid join condition",
+            ],
+        ],
+        inputs=[question, schema, correct_query, student_query, error_message],
+    )
+
+if __name__ == "__main__":
+    demo.launch()

config/codebert_labels.yaml

@@ -0,0 +1,29 @@
+# Primary labels for CodeBERT cross-encoder training
+labels:
+  - JOIN_ERROR
+  - AGGREGATION_ERROR
+  - FILTER_ERROR
+  - WINDOW_FUNCTION_ERROR
+  - SUBQUERY_ERROR
+  - NULL_HANDLING_ERROR
+  - PERFORMANCE_ERROR
+  - LOGICAL_ERROR
+  - SYNTAX_ERROR
+
+# Map dataset label_name values → one or more CodeBERT labels (multi-label)
+alias_map:
+  JOIN_ERROR: [JOIN_ERROR]
+  AGGREGATION_ERROR: [AGGREGATION_ERROR]
+  HAVING_WHERE_ERROR: [AGGREGATION_ERROR]
+  FILTERING_ERROR: [FILTER_ERROR]
+  WINDOW_FUNCTION_ERROR: [WINDOW_FUNCTION_ERROR]
+  SUBQUERY_ERROR: [SUBQUERY_ERROR]
+  NULL_HANDLING_ERROR: [NULL_HANDLING_ERROR]
+  PERFORMANCE_ERROR: [PERFORMANCE_ERROR]
+  LOGICAL_QUERY_ERROR: [LOGICAL_ERROR]
+  SYNTAX_ERROR: [SYNTAX_ERROR]
+  DATE_FUNCTION_ERROR: [SYNTAX_ERROR]
+  COLUMN_REFERENCE_ERROR: [SYNTAX_ERROR]
+  TABLE_REFERENCE_ERROR: [SYNTAX_ERROR]
+  DATA_TYPE_ERROR: [SYNTAX_ERROR]
+  DUPLICATE_RECORD_ERROR: [FILTER_ERROR]

config/error_categories.yaml

@@ -0,0 +1,46 @@
+categories:
+  - id: 0
+    name: SYNTAX_ERROR
+    description: Missing comma, bracket, quote
+  - id: 1
+    name: JOIN_ERROR
+    description: Missing ON, wrong join condition
+  - id: 2
+    name: AGGREGATION_ERROR
+    description: Missing GROUP BY
+  - id: 3
+    name: HAVING_WHERE_ERROR
+    description: Using aggregate in WHERE
+  - id: 4
+    name: SUBQUERY_ERROR
+    description: Multiple rows returned
+  - id: 5
+    name: WINDOW_FUNCTION_ERROR
+    description: Incorrect OVER/PARTITION BY
+  - id: 6
+    name: NULL_HANDLING_ERROR
+    description: "= NULL instead of IS NULL"
+  - id: 7
+    name: DATE_FUNCTION_ERROR
+    description: Incorrect date format/function
+  - id: 8
+    name: COLUMN_REFERENCE_ERROR
+    description: Column doesn't exist
+  - id: 9
+    name: TABLE_REFERENCE_ERROR
+    description: Table doesn't exist
+  - id: 10
+    name: DATA_TYPE_ERROR
+    description: Comparing integer with string
+  - id: 11
+    name: DUPLICATE_RECORD_ERROR
+    description: Missing DISTINCT
+  - id: 12
+    name: LOGICAL_QUERY_ERROR
+    description: Query runs but answer is wrong
+  - id: 13
+    name: PERFORMANCE_ERROR
+    description: "SELECT *, inefficient joins"
+  - id: 14
+    name: FILTERING_ERROR
+    description: Incorrect WHERE clause

hub/CODEBERT_MODEL_CARD.md

@@ -0,0 +1,65 @@
+---
+language: en
+license: mit
+tags:
+  - codebert
+  - sql
+  - education
+  - text-classification
+  - cross-encoder
+base_model: microsoft/codebert-base
+pipeline_tag: text-classification
+---
+
+# SQL CodeBERT Cross-Encoder
+
+Multi-label SQL error classifier using **microsoft/codebert-base** as a cross-encoder.
+
+## Input Format
+
+All fields are concatenated into one sequence:
+
+```
+QUESTION:
+{question}
+
+SCHEMA:
+{schema}
+
+STUDENT_SQL:
+{student_sql}
+
+CORRECT_SQL:
+{correct_sql}
+```
+
+## Labels
+
+`JOIN_ERROR`, `AGGREGATION_ERROR`, `FILTER_ERROR`, `WINDOW_FUNCTION_ERROR`,
+`SUBQUERY_ERROR`, `NULL_HANDLING_ERROR`, `PERFORMANCE_ERROR`, `LOGICAL_ERROR`, `SYNTAX_ERROR`
+
+## Training
+
+```bash
+python -m src.hf_train_codebert \
+  --data data/sql_errors_1m.parquet \
+  --output-dir models/codebert-cross-encoder \
+  --epochs 3 \
+  --push-to-hub \
+  --hub-model-id YOUR_USERNAME/sql-codebert-cross-encoder
+```
+
+## Inference
+
+```python
+from src.hf_predict_codebert import CodeBERTSQLErrorClassifier
+
+clf = CodeBERTSQLErrorClassifier("YOUR_USERNAME/sql-codebert-cross-encoder")
+result = clf.predict(
+    question="What is the average score per department?",
+    schema="students(id, score, department_id)",
+    student_sql="SELECT department_id, SUM(score) FROM students GROUP BY department_id",
+    correct_sql="SELECT department_id, AVG(score) FROM students GROUP BY department_id",
+)
+print(result["error_labels"])
+```

hub/MODEL_CARD.md

@@ -0,0 +1,97 @@
+---
+language: en
+license: mit
+tags:
+  - sql
+  - education
+  - text-classification
+  - sentence-transformers
+  - multi-tower
+pipeline_tag: text-classification
+---
+
+# SQL Error Classifier (Multi-Tower)
+
+Lightweight classifier that identifies **which SQL mistake area** a student is struggling with, given:
+
+- **Question** — natural-language task
+- **Schema** — available tables and columns
+- **Correct query** — reference solution
+- **Student query** — what the student submitted
+- **Error message** *(optional)* — database error text
+
+## Architecture
+
+Multi-tower semantic comparison using `sentence-transformers/all-MiniLM-L6-v2`:
+
+1. **Intent tower** — question + schema
+2. **Reference tower** — correct query
+3. **Student tower** — student query (+ error)
+4. **Comparison layer** — embedding diff, interaction, cosine similarities, SQL rule features
+5. **Linear head** — 15 error categories
+
+## Error Categories (15)
+
+| ID | Category |
+|----|----------|
+| 0 | SYNTAX_ERROR |
+| 1 | JOIN_ERROR |
+| 2 | AGGREGATION_ERROR |
+| 3 | HAVING_WHERE_ERROR |
+| 4 | SUBQUERY_ERROR |
+| 5 | WINDOW_FUNCTION_ERROR |
+| 6 | NULL_HANDLING_ERROR |
+| 7 | DATE_FUNCTION_ERROR |
+| 8 | COLUMN_REFERENCE_ERROR |
+| 9 | TABLE_REFERENCE_ERROR |
+| 10 | DATA_TYPE_ERROR |
+| 11 | DUPLICATE_RECORD_ERROR |
+| 12 | LOGICAL_QUERY_ERROR |
+| 13 | PERFORMANCE_ERROR |
+| 14 | FILTERING_ERROR |
+
+## Usage
+
+```python
+from src.huggingface import SQLLErrorClassifierHF
+
+clf = SQLLErrorClassifierHF.from_pretrained("YOUR_USERNAME/sql-error-classifier")
+
+result = clf.predict(
+    question="What is the average score of students in each department?",
+    schema="students(id, name, score, department_id) | departments(id, name)",
+    correct_query="SELECT department_id, AVG(score) FROM students GROUP BY department_id",
+    student_query="SELECT department_id, SUM(score) FROM students GROUP BY department_id",
+)
+
+print(result["label_name"])    # LOGICAL_QUERY_ERROR
+print(result["confidence"])    # 0.94
+print(result["similarities"])  # semantic alignment scores
+```
+
+## Gradio Demo
+
+Deploy as a [Hugging Face Space](https://huggingface.co/docs/hub/spaces) with `app.py` from this repository.
+
+## Model Details
+
+- **Encoder**: `sentence-transformers/all-MiniLM-L6-v2` (loaded from Hub, not bundled)
+- **Head**: scikit-learn SGDClassifier + StandardScaler
+- **Size**: ~5 MB classifier head (encoder ~80 MB, cached separately)
+- **Inference**: ~100–200 ms on CPU
+
+## Training Data
+
+Synthetically generated from exercise templates with per-category error injectors.
+1M balanced samples across 15 classes.
+
+## Citation
+
+```bibtex
+@misc{sql-error-classifier,
+  title  = {SQL Error Classifier - Multi-Tower},
+  author = {SQLErrorClassification},
+  year   = {2025},
+  publisher = {Hugging Face},
+}
+```

requirements.txt

@@ -0,0 +1,15 @@
+numpy>=1.24.0
+pandas>=2.0.0
+scikit-learn>=1.3.0
+joblib>=1.3.0
+pyarrow>=14.0.0
+tqdm>=4.66.0
+pyyaml>=6.0
+matplotlib>=3.7.0
+sentence-transformers>=2.2.0
+torch>=2.0.0
+transformers>=4.36.0
+accelerate>=0.25.0
+datasets>=2.16.0
+huggingface_hub>=0.20.0
+gradio>=4.44.0

scripts/create_hf_package.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python3
+"""Create a local Hugging Face Hub package from a trained model."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from src.huggingface import package_for_hub
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Create HF Hub package locally")
+    parser.add_argument(
+        "--model",
+        type=Path,
+        default=PROJECT_ROOT / "models" / "multi_tower_dev.joblib",
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=PROJECT_ROOT / "models" / "hf_package",
+    )
+    args = parser.parse_args()
+
+    out = package_for_hub(args.model, args.output)
+    print(f"Package ready at {out}")
+    print("Files:", [p.name for p in out.iterdir()])
+
+
+if __name__ == "__main__":
+    main()

scripts/deploy_train_space.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+# Deploy training Space to Hugging Face.
+# Usage: ./scripts/deploy_train_space.sh YOUR_HF_USERNAME/sql-error-classifier-train
+set -euo pipefail
+
+SPACE_ID="${1:-}"
+if [[ -z "${SPACE_ID}" ]]; then
+  echo "Usage: $0 YOUR_USERNAME/sql-error-classifier-train"
+  exit 1
+fi
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}"
+
+if [[ -z "${TOKEN}" ]]; then
+  echo "Set HF_TOKEN before deploying."
+  exit 1
+fi
+
+WORKDIR=$(mktemp -d)
+trap 'rm -rf "${WORKDIR}"' EXIT
+
+echo "==> Preparing Space files in ${WORKDIR}..."
+rsync -a \
+  --exclude '.venv' \
+  --exclude 'models' \
+  --exclude '__pycache__' \
+  --exclude '.git' \
+  "${ROOT}/" "${WORKDIR}/"
+
+cp "${ROOT}/README_TRAIN_SPACE.md" "${WORKDIR}/README.md"
+
+echo "==> Creating / updating Space ${SPACE_ID}..."
+python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+api.create_repo("${SPACE_ID}", repo_type="space", space_sdk="gradio", exist_ok=True)
+PY
+
+echo "==> Uploading to Hugging Face Space..."
+python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+api.upload_folder(
+    folder_path="${WORKDIR}",
+    repo_id="${SPACE_ID}",
+    repo_type="space",
+    commit_message="Deploy CodeBERT training Space",
+)
+PY
+
+echo "==> Done: https://huggingface.co/spaces/${SPACE_ID}"
+echo "Next: Space Settings → Hardware → GPU t4-small"
+echo "      Space Settings → Secrets → HF_TOKEN"

scripts/push_to_hub.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""Package and push the SQL error classifier to Hugging Face Hub."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from huggingface_hub import HfApi, create_repo
+
+from src.huggingface import package_for_hub
+DEFAULT_MODEL = PROJECT_ROOT / "models" / "multi_tower_dev.joblib"
+DEFAULT_PACKAGE = PROJECT_ROOT / "models" / "hf_package"
+MODEL_CARD = PROJECT_ROOT / "hub" / "MODEL_CARD.md"
+
+
+def push(
+    model_path: Path = DEFAULT_MODEL,
+    package_dir: Path = DEFAULT_PACKAGE,
+    repo_id: str = "",
+    private: bool = False,
+    token: str | None = None,
+) -> str:
+    if not repo_id:
+        raise ValueError("--repo-id is required (e.g. your-username/sql-error-classifier)")
+
+    token = token or os.getenv("HF_TOKEN")
+    api = HfApi(token=token)
+
+    print(f"Packaging model from {model_path}...")
+    package_for_hub(model_path, package_dir)
+
+    print(f"Creating repo {repo_id}...")
+    create_repo(repo_id, repo_type="model", private=private, exist_ok=True, token=token)
+
+    print("Uploading model files...")
+    api.upload_folder(
+        folder_path=str(package_dir),
+        repo_id=repo_id,
+        repo_type="model",
+        token=token,
+    )
+
+    if MODEL_CARD.exists():
+        api.upload_file(
+            path_or_fileobj=str(MODEL_CARD),
+            path_in_repo="README.md",
+            repo_id=repo_id,
+            repo_type="model",
+            token=token,
+        )
+
+    url = f"https://huggingface.co/{repo_id}"
+    print(f"Done: {url}")
+    return url
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Push SQL error classifier to HF Hub")
+    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
+    parser.add_argument("--package-dir", type=Path, default=DEFAULT_PACKAGE)
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        required=True,
+        help="Hugging Face repo id, e.g. nishantgupta/sql-error-classifier",
+    )
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument("--token", type=str, default=None)
+    args = parser.parse_args()
+
+    push(
+        model_path=args.model,
+        package_dir=args.package_dir,
+        repo_id=args.repo_id,
+        private=args.private,
+        token=args.token,
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/run_codebert_training.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+DATA="${1:-data/sql_errors_dev.parquet}"
+OUTPUT="${2:-models/codebert-cross-encoder}"
+SAMPLES="${3:-}"
+
+CMD=(python -m src.hf_train_codebert --data "${DATA}" --output-dir "${OUTPUT}")
+
+if [[ -n "${SAMPLES}" ]]; then
+  CMD+=(--max-samples "${SAMPLES}")
+fi
+
+echo "==> Training CodeBERT cross-encoder..."
+"${CMD[@]}"
+
+echo "==> Done. Model at ${OUTPUT}"

scripts/run_pipeline.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SAMPLES="${1:-1000000}"
+WORKERS="${2:-8}"
+
+echo "==> Generating ${SAMPLES} labeled SQL samples..."
+python -m src.generate_dataset --samples "${SAMPLES}" --workers "${WORKERS}"
+
+echo "==> Training classifier..."
+python -m src.train
+
+echo "==> Evaluating..."
+python -m src.evaluate
+
+echo "==> Done. Model at models/sql_error_classifier.joblib"

src/__init__.py

@@ -0,0 +1 @@
+"""SQL error classification package."""

src/categories.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List
+
+import yaml
+
+CONFIG_PATH = Path(__file__).resolve().parent.parent / "config" / "error_categories.yaml"
+
+
+@dataclass(frozen=True)
+class ErrorCategory:
+    id: int
+    name: str
+    description: str
+
+
+def load_categories(config_path: Path = CONFIG_PATH) -> List[ErrorCategory]:
+    with open(config_path) as f:
+        data = yaml.safe_load(f)
+    return [
+        ErrorCategory(id=c["id"], name=c["name"], description=c["description"])
+        for c in data["categories"]
+    ]
+
+
+def id_to_name(categories: List[ErrorCategory] | None = None) -> Dict[int, str]:
+    cats = categories or load_categories()
+    return {c.id: c.name for c in cats}
+
+
+def name_to_id(categories: List[ErrorCategory] | None = None) -> Dict[str, int]:
+    cats = categories or load_categories()
+    return {c.name: c.id for c in cats}

src/codebert_dataset.py

@@ -0,0 +1,117 @@
+"""PyTorch Dataset and preprocessing for CodeBERT Trainer."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+from transformers import PreTrainedTokenizerBase
+
+from src.codebert_formatting import format_cross_encoder_input
+from src.codebert_labels import label_to_multihot, load_codebert_labels
+
+
+def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
+    """Map project column names to the canonical training schema."""
+    col_map = {
+        "query": "student_sql",
+        "correct_query": "correct_sql",
+        "label_name": "error_labels",
+    }
+    out = df.rename(columns={k: v for k, v in col_map.items() if k in df.columns}).copy()
+
+    required = ["question", "schema", "student_sql", "correct_sql", "error_labels"]
+    missing = [c for c in required if c not in out.columns]
+    if missing:
+        raise ValueError(
+            f"Dataset missing required columns: {missing}. "
+            f"Expected {required} (or aliases query/correct_query/label_name)."
+        )
+    return out
+
+
+class SQLCodeBERTDataset(Dataset):
+    """Tokenized SQL error dataset for Hugging Face Trainer."""
+
+    def __init__(
+        self,
+        df: pd.DataFrame,
+        tokenizer: PreTrainedTokenizerBase,
+        label_list: Optional[List[str]] = None,
+        max_length: int = 512,
+    ):
+        self.df = normalize_dataframe(df).reset_index(drop=True)
+        self.tokenizer = tokenizer
+        self.label_list = label_list or load_codebert_labels()
+        self.max_length = max_length
+        self.num_labels = len(self.label_list)
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        row = self.df.iloc[idx]
+        text = format_cross_encoder_input(
+            question=str(row["question"]),
+            schema=str(row["schema"]),
+            student_sql=str(row["student_sql"]),
+            correct_sql=str(row["correct_sql"]),
+        )
+        encoded = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            padding=False,
+            return_tensors=None,
+        )
+        labels = label_to_multihot(str(row["error_labels"]), self.label_list)
+        encoded["labels"] = labels.tolist()
+        return encoded
+
+
+class SQLCodeBERTDataCollator:
+    """Pad batches dynamically for Trainer."""
+
+    def __init__(self, tokenizer: PreTrainedTokenizerBase):
+        self.tokenizer = tokenizer
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        labels = [f.pop("labels") for f in features]
+        batch = self.tokenizer.pad(features, padding=True, return_tensors="pt")
+        batch["labels"] = torch.tensor(labels, dtype=torch.float)
+        return batch
+
+
+def prepare_datasets(
+    df: pd.DataFrame,
+    tokenizer: PreTrainedTokenizerBase,
+    test_size: float = 0.1,
+    val_size: float = 0.1,
+    max_length: int = 512,
+    seed: int = 42,
+) -> tuple[SQLCodeBERTDataset, SQLCodeBERTDataset, SQLCodeBERTDataset]:
+    from sklearn.model_selection import train_test_split
+
+    df = normalize_dataframe(df)
+    trainval, test_df = train_test_split(
+        df,
+        test_size=test_size,
+        random_state=seed,
+        stratify=df["error_labels"],
+    )
+    relative_val = val_size / (1 - test_size)
+    train_df, val_df = train_test_split(
+        trainval,
+        test_size=relative_val,
+        random_state=seed,
+        stratify=trainval["error_labels"],
+    )
+
+    return (
+        SQLCodeBERTDataset(train_df, tokenizer, max_length=max_length),
+        SQLCodeBERTDataset(val_df, tokenizer, max_length=max_length),
+        SQLCodeBERTDataset(test_df, tokenizer, max_length=max_length),
+    )

src/codebert_formatting.py

@@ -0,0 +1,28 @@
+"""Cross-encoder input formatting for CodeBERT."""
+
+from __future__ import annotations
+
+QUESTION_TAG = "QUESTION:"
+SCHEMA_TAG = "SCHEMA:"
+STUDENT_TAG = "STUDENT_SQL:"
+CORRECT_TAG = "CORRECT_SQL:"
+
+
+def format_cross_encoder_input(
+    question: str,
+    schema: str,
+    student_sql: str,
+    correct_sql: str,
+) -> str:
+    """
+    Concatenate all fields into a single CodeBERT input sequence.
+
+    The model attends jointly across question intent, schema, student SQL,
+    and the reference solution — cross-encoder style in one forward pass.
+    """
+    return (
+        f"{QUESTION_TAG}\n{question.strip()}\n\n"
+        f"{SCHEMA_TAG}\n{schema.strip()}\n\n"
+        f"{STUDENT_TAG}\n{student_sql.strip()}\n\n"
+        f"{CORRECT_TAG}\n{correct_sql.strip()}"
+    )

src/codebert_labels.py

@@ -0,0 +1,82 @@
+"""Label utilities for CodeBERT multi-label classification."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Dict, List, Sequence, Union
+
+import numpy as np
+import yaml
+
+CONFIG_PATH = (
+    Path(__file__).resolve().parent.parent / "config" / "codebert_labels.yaml"
+)
+
+
+def load_codebert_labels(config_path: Path = CONFIG_PATH) -> List[str]:
+    with open(config_path) as f:
+        data = yaml.safe_load(f)
+    return list(data["labels"])
+
+
+def load_alias_map(config_path: Path = CONFIG_PATH) -> Dict[str, List[str]]:
+    with open(config_path) as f:
+        data = yaml.safe_load(f)
+    return {k: list(v) for k, v in data["alias_map"].items()}
+
+
+def label_to_multihot(
+    error_labels: Union[str, Sequence[str]],
+    label_list: List[str] | None = None,
+    alias_map: Dict[str, List[str]] | None = None,
+) -> np.ndarray:
+    """
+    Convert error label(s) to multi-hot vector.
+
+    Accepts:
+      - comma-separated string: "JOIN_ERROR,AGGREGATION_ERROR"
+      - list of label strings
+      - single dataset label_name (resolved via alias_map)
+    """
+    labels = label_list or load_codebert_labels()
+    aliases = alias_map or load_alias_map()
+    index = {name: i for i, name in enumerate(labels)}
+    vec = np.zeros(len(labels), dtype=np.float32)
+
+    if isinstance(error_labels, str):
+        raw = [s.strip() for s in error_labels.split(",") if s.strip()]
+        if len(raw) == 1 and raw[0] in aliases:
+            raw = aliases[raw[0]]
+        elif len(raw) == 1 and raw[0] in index:
+            raw = [raw[0]]
+        elif len(raw) == 1 and raw[0] not in index:
+            mapped = aliases.get(raw[0], [])
+            raw = mapped
+    else:
+        raw = list(error_labels)
+        expanded: List[str] = []
+        for item in raw:
+            if item in aliases:
+                expanded.extend(aliases[item])
+            elif item in index:
+                expanded.append(item)
+        raw = expanded
+
+    for name in raw:
+        if name not in index:
+            raise ValueError(f"Unknown label '{name}'. Expected one of {labels}")
+        vec[index[name]] = 1.0
+
+    if vec.sum() == 0:
+        raise ValueError(f"No valid labels found in {error_labels}")
+    return vec
+
+
+def multihot_to_label_names(
+    vec: np.ndarray,
+    label_list: List[str] | None = None,
+    threshold: float = 0.5,
+) -> List[str]:
+    labels = label_list or load_codebert_labels()
+    indices = np.where(vec >= threshold)[0]
+    return [labels[i] for i in indices]

src/cross_encoder_model.py

@@ -0,0 +1,312 @@
+"""Cross-encoder architecture for SQL error classification."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Optional, Tuple
+
+import numpy as np
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import StandardScaler
+
+from src.multi_tower_model import QueryContext, contexts_from_dataframe
+from src.sql_features import extract_sql_features
+
+DEFAULT_CROSS_ENCODER = "cross-encoder/ms-marco-MiniLM-L6-v2"
+DEFAULT_FINETUNED_CE = "cross-encoder/ms-marco-MiniLM-L6-v2"
+
+PAIR_NAMES = (
+    "intent_vs_student",
+    "reference_vs_student",
+    "intent_vs_reference",
+)
+
+
+@dataclass(frozen=True)
+class CrossEncoderPair:
+    name: str
+    text_a: str
+    text_b: str
+
+
+def _intent_text(ctx: QueryContext) -> str:
+    return f"QUESTION: {ctx.question} SCHEMA: {ctx.schema}"
+
+
+def _reference_text(ctx: QueryContext) -> str:
+    return f"REFERENCE: {ctx.correct_query}"
+
+
+def _student_text(ctx: QueryContext) -> str:
+    parts = [f"STUDENT: {ctx.student_query}"]
+    if ctx.error_message:
+        parts.append(f"ERROR: {ctx.error_message}")
+    return " ".join(parts)
+
+
+def _context_text(ctx: QueryContext) -> str:
+    """Full task context for fine-tuned cross-encoder."""
+    return (
+        f"QUESTION: {ctx.question} "
+        f"SCHEMA: {ctx.schema} "
+        f"REFERENCE: {ctx.correct_query}"
+    )
+
+
+def build_pairs(ctx: QueryContext) -> List[CrossEncoderPair]:
+    intent, reference, student = (
+        _intent_text(ctx),
+        _reference_text(ctx),
+        _student_text(ctx),
+    )
+    return [
+        CrossEncoderPair("intent_vs_student", intent, student),
+        CrossEncoderPair("reference_vs_student", reference, student),
+        CrossEncoderPair("intent_vs_reference", intent, reference),
+    ]
+
+
+class CrossEncoderClassifier:
+    """
+    Hybrid cross-encoder: frozen pairwise relevance + linear head.
+
+    Unlike bi-encoders (multi-tower), the cross-encoder attends jointly over
+    each (context, student) pair — better for logical and filtering errors.
+
+    Three pairs are scored:
+      1. intent vs student     — does the query address the question?
+      2. reference vs student  — how far is the student from the answer?
+      3. intent vs reference   — task-answer alignment baseline
+
+    Pair scores + SQL rule features → LogisticRegression → 15 classes.
+    """
+
+    def __init__(
+        self,
+        cross_encoder_name: str = DEFAULT_CROSS_ENCODER,
+        batch_size: int = 32,
+        max_length: int = 512,
+    ):
+        self.cross_encoder_name = cross_encoder_name
+        self.batch_size = batch_size
+        self.max_length = max_length
+        self.cross_encoder = None
+        self.scaler = StandardScaler()
+        self.clf = LogisticRegression(
+            max_iter=1000,
+            solver="lbfgs",
+            class_weight="balanced",
+            random_state=42,
+        )
+        self.classes_: Optional[np.ndarray] = None
+
+    def _load_cross_encoder(self):
+        if self.cross_encoder is None:
+            from sentence_transformers import CrossEncoder
+
+            self.cross_encoder = CrossEncoder(
+                self.cross_encoder_name,
+                max_length=self.max_length,
+            )
+
+    def _pair_batches(self, contexts: List[QueryContext]) -> List[List[Tuple[str, str]]]:
+        """One batch list per pair type across all contexts."""
+        pair_lists: List[List[Tuple[str, str]]] = [[], [], []]
+        for ctx in contexts:
+            pairs = build_pairs(ctx)
+            for i, pair in enumerate(pairs):
+                pair_lists[i].append((pair.text_a, pair.text_b))
+        return pair_lists
+
+    def _score_pairs(
+        self,
+        contexts: List[QueryContext],
+        show_progress: bool = False,
+    ) -> np.ndarray:
+        self._load_cross_encoder()
+        pair_batches = self._pair_batches(contexts)
+        scores = []
+        for batch in pair_batches:
+            raw = self.cross_encoder.predict(
+                batch,
+                batch_size=self.batch_size,
+                show_progress_bar=show_progress,
+            )
+            scores.append(np.asarray(raw, dtype=np.float64).reshape(-1, 1))
+        return np.hstack(scores)  # (n, 3)
+
+    def _build_features(
+        self,
+        contexts: List[QueryContext],
+        show_progress: bool = False,
+    ) -> np.ndarray:
+        pair_scores = self._score_pairs(contexts, show_progress=show_progress)
+        s_is, s_rs, s_ir = pair_scores[:, 0], pair_scores[:, 1], pair_scores[:, 2]
+
+        derived = np.column_stack(
+            [
+                s_rs - s_is,          # reference closer than intent?
+                s_is - s_ir,          # student-intent gap vs baseline
+                s_rs - s_ir,          # student-reference gap vs baseline
+                s_is * s_rs,          # interaction
+                np.abs(s_rs - s_is),  # intent-reference disagreement
+            ]
+        )
+
+        sql_feats = np.array(
+            [extract_sql_features(c.student_query, c.correct_query) for c in contexts],
+            dtype=np.float64,
+        )
+
+        return np.hstack([pair_scores, derived, sql_feats])
+
+    def _prepare_features(self, contexts: List[QueryContext]) -> np.ndarray:
+        X = self.scaler.transform(self._build_features(contexts))
+        return np.nan_to_num(X, nan=0.0, posinf=1e3, neginf=-1e3)
+
+    def fit(self, contexts: List[QueryContext], y: np.ndarray) -> "CrossEncoderClassifier":
+        X = self._build_features(contexts, show_progress=True)
+        X = self.scaler.fit_transform(X)
+        X = np.nan_to_num(X, nan=0.0, posinf=1e3, neginf=-1e3)
+        self.clf.fit(X, y)
+        self.classes_ = self.clf.classes_
+        return self
+
+    def predict(self, contexts: List[QueryContext]) -> np.ndarray:
+        return self.clf.predict(self._prepare_features(contexts))
+
+    def predict_proba(self, contexts: List[QueryContext]) -> np.ndarray:
+        return self.clf.predict_proba(self._prepare_features(contexts))
+
+    def explain_pair_scores(self, ctx: QueryContext) -> dict:
+        scores = self._score_pairs([ctx])[0]
+        return {
+            PAIR_NAMES[0]: float(scores[0]),
+            PAIR_NAMES[1]: float(scores[1]),
+            PAIR_NAMES[2]: float(scores[2]),
+        }
+
+
+class FineTunedCrossEncoderClassifier:
+    """
+    End-to-end fine-tuned cross-encoder (highest accuracy).
+
+    Single cross-attention pass over [task_context | student_query] with
+    num_labels=15. Slower to train; best on smaller high-quality datasets.
+    """
+
+    def __init__(
+        self,
+        cross_encoder_name: str = DEFAULT_FINETUNED_CE,
+        batch_size: int = 16,
+        max_length: int = 512,
+        num_labels: int = 15,
+    ):
+        self.cross_encoder_name = cross_encoder_name
+        self.batch_size = batch_size
+        self.max_length = max_length
+        self.num_labels = num_labels
+        self.model = None
+        self.classes_: Optional[np.ndarray] = None
+
+    def _load_model(self, num_labels: Optional[int] = None):
+        if self.model is None:
+            from sentence_transformers import CrossEncoder
+
+            self.model = CrossEncoder(
+                self.cross_encoder_name,
+                num_labels=num_labels or self.num_labels,
+                max_length=self.max_length,
+            )
+
+    def _to_examples(self, contexts: List[QueryContext], labels: Optional[np.ndarray] = None):
+        from sentence_transformers import InputExample
+
+        examples = []
+        for i, ctx in enumerate(contexts):
+            label = float(labels[i]) if labels is not None else 0.0
+            examples.append(
+                InputExample(
+                    texts=[_context_text(ctx), _student_text(ctx)],
+                    label=label,
+                )
+            )
+        return examples
+
+    def fit(
+        self,
+        contexts: List[QueryContext],
+        y: np.ndarray,
+        epochs: int = 1,
+        warmup_steps: int = 100,
+        output_path: Optional[Path] = None,
+    ) -> "FineTunedCrossEncoderClassifier":
+        from torch.utils.data import DataLoader
+
+        self._load_model(num_labels=len(np.unique(y)))
+        train_examples = self._to_examples(contexts, y)
+        loader = DataLoader(
+            train_examples,
+            shuffle=True,
+            batch_size=self.batch_size,
+        )
+        self.model.fit(
+            train_dataloader=loader,
+            epochs=epochs,
+            warmup_steps=min(warmup_steps, max(10, len(train_examples) // 10)),
+            show_progress_bar=True,
+            output_path=str(output_path) if output_path else None,
+        )
+        self.classes_ = np.sort(np.unique(y))
+        return self
+
+    def predict(self, contexts: List[QueryContext]) -> np.ndarray:
+        self._load_model()
+        pairs = [[_context_text(c), _student_text(c)] for c in contexts]
+        logits = self.model.predict(
+            pairs,
+            batch_size=self.batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+        )
+        logits = np.asarray(logits)
+        if logits.ndim == 1:
+            return logits.astype(int)
+        return logits.argmax(axis=1)
+
+    def predict_proba(self, contexts: List[QueryContext]) -> np.ndarray:
+        self._load_model()
+        pairs = [[_context_text(c), _student_text(c)] for c in contexts]
+        logits = self.model.predict(
+            pairs,
+            batch_size=self.batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+        )
+        logits = np.asarray(logits, dtype=np.float64)
+        if logits.ndim == 1:
+            # binary fallback
+            probs = np.zeros((len(contexts), len(self.classes_)))
+            for i, pred in enumerate(logits.astype(int)):
+                idx = np.where(self.classes_ == pred)[0][0]
+                probs[i, idx] = 1.0
+            return probs
+        # softmax
+        exp = np.exp(logits - logits.max(axis=1, keepdims=True))
+        return exp / exp.sum(axis=1, keepdims=True)
+
+    def save(self, path: Path) -> Path:
+        path.mkdir(parents=True, exist_ok=True)
+        self._load_model()
+        self.model.save(str(path))
+        return path
+
+    @classmethod
+    def load(cls, path: Path) -> "FineTunedCrossEncoderClassifier":
+        from sentence_transformers import CrossEncoder
+
+        instance = cls()
+        instance.model = CrossEncoder(str(path))
+        instance.classes_ = np.arange(instance.model.num_labels)
+        return instance

src/evaluate.py

@@ -0,0 +1,114 @@
+"""Evaluate trained model with confusion matrix and per-class metrics."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from sklearn.metrics import ConfusionMatrixDisplay, classification_report
+
+from src.categories import id_to_name, load_categories
+from src.model import DEFAULT_MODEL_PATH, combine_features, load_model
+from src.cross_encoder_model import (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+)
+from src.multi_tower_model import MultiTowerClassifier, contexts_from_dataframe
+
+CONTEXT_MODELS = (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+    MultiTowerClassifier,
+)
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_DATA = PROJECT_ROOT / "data" / "sql_errors_1m.parquet"
+DEFAULT_OUTPUT = PROJECT_ROOT / "models" / "evaluation"
+
+
+def evaluate(
+    data_path: Path = DEFAULT_DATA,
+    model_path: Path = DEFAULT_MODEL_PATH,
+    output_dir: Path = DEFAULT_OUTPUT,
+    sample_size: int = 100_000,
+    use_error_message: bool = True,
+    seed: int = 42,
+) -> dict:
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    df = pd.read_parquet(data_path)
+    if len(df) > sample_size:
+        df = df.sample(n=sample_size, random_state=seed)
+
+    labels = df["label_id"].values
+    model = load_model(model_path)
+
+    if isinstance(model, CONTEXT_MODELS):
+        if not use_error_message and "error_message" in df.columns:
+            df = df.drop(columns=["error_message"])
+        preds = model.predict(contexts_from_dataframe(df))
+    else:
+        texts = combine_features(
+            queries=df["query"].tolist(),
+            error_messages=df["error_message"].tolist() if use_error_message else None,
+            schemas=df["schema"].tolist() if "schema" in df.columns else None,
+            questions=df["question"].tolist() if "question" in df.columns else None,
+        )
+        preds = model.predict(texts)
+
+    categories = load_categories()
+    target_names = [c.name for c in categories]
+    report = classification_report(
+        labels, preds, target_names=target_names, output_dict=True, zero_division=0
+    )
+
+    with open(output_dir / "classification_report.json", "w") as f:
+        json.dump(report, f, indent=2)
+
+    cm = ConfusionMatrixDisplay.from_predictions(
+        labels,
+        preds,
+        display_labels=target_names,
+        xticks_rotation=90,
+        cmap="Blues",
+        colorbar=False,
+    )
+    fig = cm.figure_
+    fig.set_size_inches(14, 12)
+    fig.tight_layout()
+    fig.savefig(output_dir / "confusion_matrix.png", dpi=150)
+    plt.close(fig)
+
+    print(f"Accuracy: {report['accuracy']:.4f}")
+    print(f"Reports saved to {output_dir}")
+    return report
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate SQL error classifier")
+    parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
+    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL_PATH)
+    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    parser.add_argument("--sample-size", type=int, default=100_000)
+    parser.add_argument("--no-error-message", action="store_true")
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    evaluate(
+        data_path=args.data,
+        model_path=args.model,
+        output_dir=args.output,
+        sample_size=args.sample_size,
+        use_error_message=not args.no_error_message,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/exercises.py

@@ -0,0 +1,228 @@
+"""Generate playground exercises: schema + question + correct SQL."""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import Callable, Dict, List
+
+
+@dataclass(frozen=True)
+class Exercise:
+    schema: str
+    question: str
+    correct_query: str
+    tables: tuple[str, ...]
+    columns: tuple[str, ...]
+
+
+def _fmt_schema(tables: Dict[str, List[str]]) -> str:
+    parts = [f"{name}({', '.join(cols)})" for name, cols in tables.items()]
+    return " | ".join(parts)
+
+
+ExerciseBuilder = Callable[[random.Random], Exercise]
+
+EXERCISE_BUILDERS: List[ExerciseBuilder] = []
+
+
+def _register(builder: ExerciseBuilder) -> ExerciseBuilder:
+    EXERCISE_BUILDERS.append(builder)
+    return builder
+
+
+@_register
+def exercise_avg_by_department(rng: random.Random) -> Exercise:
+    tables = {
+        "students": ["id", "name", "email", "score", "department_id"],
+        "departments": ["id", "name", "city"],
+    }
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="What is the average score of students in each department?",
+        correct_query=(
+            "SELECT department_id, AVG(score) "
+            "FROM students GROUP BY department_id"
+        ),
+        tables=tuple(tables),
+        columns=("department_id", "score"),
+    )
+
+
+@_register
+def exercise_student_department_names(rng: random.Random) -> Exercise:
+    tables = {
+        "students": ["id", "name", "department_id"],
+        "departments": ["id", "name"],
+    }
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="List each student's name along with their department name.",
+        correct_query=(
+            "SELECT students.name, departments.name "
+            "FROM students "
+            "INNER JOIN departments ON students.department_id = departments.id"
+        ),
+        tables=tuple(tables),
+        columns=("name", "department_id"),
+    )
+
+
+@_register
+def exercise_high_scoring_students(rng: random.Random) -> Exercise:
+    threshold = rng.randint(70, 90)
+    tables = {"students": ["id", "name", "age", "score", "status"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question=(
+            f"Find names of students older than 18 with a score above {threshold}."
+        ),
+        correct_query=(
+            f"SELECT name FROM students "
+            f"WHERE age > 18 AND score > {threshold}"
+        ),
+        tables=tuple(tables),
+        columns=("name", "age", "score", "status"),
+    )
+
+
+@_register
+def exercise_unique_cities(rng: random.Random) -> Exercise:
+    tables = {"students": ["id", "name", "city", "country"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="List the unique cities where students live.",
+        correct_query="SELECT DISTINCT city FROM students",
+        tables=tuple(tables),
+        columns=("city",),
+    )
+
+
+@_register
+def exercise_top_scorer(rng: random.Random) -> Exercise:
+    tables = {"students": ["id", "name", "score"], "grades": ["id", "score"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="Find students whose score equals the highest score in the class.",
+        correct_query=(
+            "SELECT name FROM students "
+            "WHERE score = (SELECT MAX(score) FROM grades)"
+        ),
+        tables=tuple(tables),
+        columns=("name", "score"),
+    )
+
+
+@_register
+def exercise_departments_over_budget(rng: random.Random) -> Exercise:
+    budget = rng.randint(3, 8)
+    tables = {"employees": ["id", "name", "department_id", "salary"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question=f"Which departments have more than {budget} employees?",
+        correct_query=(
+            f"SELECT department_id, COUNT(*) AS cnt "
+            f"FROM employees GROUP BY department_id "
+            f"HAVING COUNT(*) > {budget}"
+        ),
+        tables=tuple(tables),
+        columns=("department_id", "salary"),
+    )
+
+
+@_register
+def exercise_recent_orders(rng: random.Random) -> Exercise:
+    year = rng.randint(2020, 2024)
+    tables = {"orders": ["id", "customer_id", "amount", "order_date", "status"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question=f"Show orders placed on or after January 1, {year}.",
+        correct_query=(
+            f"SELECT id, amount FROM orders "
+            f"WHERE order_date >= DATE '{year}-01-01'"
+        ),
+        tables=tuple(tables),
+        columns=("order_date", "amount", "status"),
+    )
+
+
+@_register
+def exercise_missing_email(rng: random.Random) -> Exercise:
+    tables = {"students": ["id", "name", "email", "phone"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="Find students who have not provided an email address.",
+        correct_query="SELECT name FROM students WHERE email IS NULL",
+        tables=tuple(tables),
+        columns=("email", "name"),
+    )
+
+
+@_register
+def exercise_rank_by_score(rng: random.Random) -> Exercise:
+    tables = {"students": ["id", "name", "score", "department_id"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="Rank students by score within each department.",
+        correct_query=(
+            "SELECT name, score, "
+            "RANK() OVER (PARTITION BY department_id ORDER BY score DESC) AS rnk "
+            "FROM students"
+        ),
+        tables=tuple(tables),
+        columns=("name", "score", "department_id"),
+    )
+
+
+@_register
+def exercise_course_enrollment_count(rng: random.Random) -> Exercise:
+    tables = {
+        "courses": ["id", "title"],
+        "enrollments": ["id", "course_id", "student_id"],
+    }
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="How many students are enrolled in each course?",
+        correct_query=(
+            "SELECT courses.title, COUNT(enrollments.student_id) AS enrolled "
+            "FROM courses "
+            "INNER JOIN enrollments ON courses.id = enrollments.course_id "
+            "GROUP BY courses.title"
+        ),
+        tables=tuple(tables),
+        columns=("title", "student_id", "course_id"),
+    )
+
+
+@_register
+def exercise_active_employees(rng: random.Random) -> Exercise:
+    tables = {"employees": ["id", "name", "salary", "status", "hire_date"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question="What is the total salary paid to active employees?",
+        correct_query=(
+            "SELECT SUM(salary) FROM employees WHERE status = 'active'"
+        ),
+        tables=tuple(tables),
+        columns=("salary", "status"),
+    )
+
+
+@_register
+def exercise_product_price_filter(rng: random.Random) -> Exercise:
+    lo, hi = rng.randint(10, 50), rng.randint(100, 500)
+    tables = {"products": ["id", "name", "price", "category"]}
+    return Exercise(
+        schema=_fmt_schema(tables),
+        question=f"List products priced between {lo} and {hi}.",
+        correct_query=(
+            f"SELECT name, price FROM products "
+            f"WHERE price BETWEEN {lo} AND {hi}"
+        ),
+        tables=tuple(tables),
+        columns=("name", "price", "category"),
+    )
+
+
+def generate_exercise(rng: random.Random) -> Exercise:
+    return rng.choice(EXERCISE_BUILDERS)(rng)

src/generate_dataset.py

@@ -0,0 +1,115 @@
+"""Generate labeled SQL error dataset at scale."""
+
+from __future__ import annotations
+
+import argparse
+import random
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import pandas as pd
+from tqdm import tqdm
+
+from src.categories import load_categories
+from src.exercises import generate_exercise
+from src.sql_templates import ERROR_INJECTORS
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "sql_errors_1m.parquet"
+
+
+def generate_dataset(
+    total_samples: int = 1_000_000,
+    output_path: Path = DEFAULT_OUTPUT,
+    batch_size: int = 10_000,
+    workers: int = 8,
+    seed: int = 42,
+) -> Path:
+    categories = load_categories()
+    label_ids = [c.id for c in categories]
+    samples_per_class = total_samples // len(label_ids)
+    remainder = total_samples % len(label_ids)
+
+    # Balanced label schedule: each class gets equal share (+1 for first `remainder` classes)
+    schedule: List[int] = []
+    for cat in categories:
+        count = samples_per_class + (1 if cat.id < remainder else 0)
+        schedule.extend([cat.id] * count)
+    random.Random(seed).shuffle(schedule)
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    chunks: List[pd.DataFrame] = []
+    num_batches = (total_samples + batch_size - 1) // batch_size
+
+    with ProcessPoolExecutor(max_workers=workers) as executor:
+        futures = []
+        offset = 0
+        for batch_idx in range(num_batches):
+            current_batch = min(batch_size, total_samples - offset)
+            batch_labels = schedule[offset : offset + current_batch]
+            futures.append(
+                executor.submit(
+                    _generate_batch_with_labels,
+                    batch_labels,
+                    seed + batch_idx,
+                )
+            )
+            offset += current_batch
+
+        for future in tqdm(as_completed(futures), total=len(futures), desc="Generating"):
+            rows = future.result()
+            chunks.append(pd.DataFrame(rows))
+
+    df = pd.concat(chunks, ignore_index=True)
+    df = df.sample(frac=1, random_state=seed).reset_index(drop=True)
+    df.to_parquet(output_path, index=False)
+
+    print(f"Saved {len(df):,} samples to {output_path}")
+    print("\nClass distribution:")
+    print(df["label_name"].value_counts().sort_index().to_string())
+    return output_path
+
+
+def _generate_batch_with_labels(label_ids: List[int], seed: int) -> List[Dict]:
+    rng = random.Random(seed)
+    categories = load_categories()
+    rows = []
+    for label_id in label_ids:
+        exercise = generate_exercise(rng)
+        injector = ERROR_INJECTORS[label_id]
+        query, error_message = injector(rng, exercise)
+        rows.append(
+            {
+                "schema": exercise.schema,
+                "question": exercise.question,
+                "correct_query": exercise.correct_query,
+                "query": query.strip(),
+                "error_message": error_message,
+                "label_id": label_id,
+                "label_name": categories[label_id].name,
+            }
+        )
+    return rows
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Generate labeled SQL error dataset")
+    parser.add_argument("--samples", type=int, default=1_000_000, help="Total samples")
+    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
+    parser.add_argument("--batch-size", type=int, default=10_000)
+    parser.add_argument("--workers", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    generate_dataset(
+        total_samples=args.samples,
+        output_path=args.output,
+        batch_size=args.batch_size,
+        workers=args.workers,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/hf_eval_codebert.py

@@ -0,0 +1,69 @@
+"""Evaluate a trained CodeBERT cross-encoder."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import pandas as pd
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer
+
+from src.codebert_dataset import SQLCodeBERTDataCollator, SQLCodeBERTDataset, normalize_dataframe
+from src.hf_metrics import build_compute_metrics, compute_multilabel_metrics
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_DATA = PROJECT_ROOT / "data" / "sql_errors_dev.parquet"
+DEFAULT_MODEL = PROJECT_ROOT / "models" / "codebert-cross-encoder"
+
+
+def evaluate(
+    model_dir: Path = DEFAULT_MODEL,
+    data_path: Path = DEFAULT_DATA,
+    sample_size: int = 10_000,
+    threshold: float = 0.5,
+    seed: int = 42,
+) -> dict:
+    df = normalize_dataframe(pd.read_parquet(data_path))
+    if len(df) > sample_size:
+        df = df.sample(n=sample_size, random_state=seed)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+    dataset = SQLCodeBERTDataset(df, tokenizer)
+
+    trainer_kwargs = dict(
+        model=model,
+        data_collator=SQLCodeBERTDataCollator(tokenizer),
+        compute_metrics=build_compute_metrics(threshold=threshold),
+    )
+    try:
+        trainer = Trainer(processing_class=tokenizer, **trainer_kwargs)
+    except TypeError:
+        trainer = Trainer(tokenizer=tokenizer, **trainer_kwargs)
+
+    output = trainer.predict(dataset)
+    metrics = compute_multilabel_metrics(
+        output.predictions, output.label_ids, threshold=threshold
+    )
+    print(json.dumps(metrics, indent=2))
+    return metrics
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate CodeBERT SQL classifier")
+    parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL)
+    parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
+    parser.add_argument("--sample-size", type=int, default=10_000)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    args = parser.parse_args()
+    evaluate(
+        model_dir=args.model_dir,
+        data_path=args.data,
+        sample_size=args.sample_size,
+        threshold=args.threshold,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/hf_metrics.py

@@ -0,0 +1,52 @@
+"""Evaluation metrics for multi-label SQL error classification."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+import numpy as np
+from sklearn.metrics import (
+    accuracy_score,
+    f1_score,
+    hamming_loss,
+    precision_score,
+    recall_score,
+)
+
+
+def sigmoid(x: np.ndarray) -> np.ndarray:
+    return 1.0 / (1.0 + np.exp(-x))
+
+
+def compute_multilabel_metrics(
+    logits: np.ndarray,
+    labels: np.ndarray,
+    threshold: float = 0.5,
+) -> Dict[str, float]:
+    probs = sigmoid(logits)
+    preds = (probs >= threshold).astype(int)
+    labels = labels.astype(int)
+
+    return {
+        "accuracy": float(accuracy_score(labels, preds)),
+        "f1_macro": float(f1_score(labels, preds, average="macro", zero_division=0)),
+        "f1_micro": float(f1_score(labels, preds, average="micro", zero_division=0)),
+        "precision_macro": float(
+            precision_score(labels, preds, average="macro", zero_division=0)
+        ),
+        "recall_macro": float(
+            recall_score(labels, preds, average="macro", zero_division=0)
+        ),
+        "hamming_loss": float(hamming_loss(labels, preds)),
+        "subset_accuracy": float((preds == labels).all(axis=1).mean()),
+    }
+
+
+def build_compute_metrics(threshold: float = 0.5):
+    """Factory for Hugging Face Trainer compute_metrics callback."""
+
+    def compute_metrics(eval_pred) -> Dict[str, float]:
+        logits, labels = eval_pred
+        return compute_multilabel_metrics(logits, labels, threshold=threshold)
+
+    return compute_metrics

src/hf_predict_codebert.py

@@ -0,0 +1,161 @@
+"""Inference for CodeBERT SQL error cross-encoder."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from src.codebert_formatting import format_cross_encoder_input
+from src.codebert_labels import load_codebert_labels, multihot_to_label_names
+from src.hf_metrics import sigmoid
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_MODEL_DIR = PROJECT_ROOT / "models" / "codebert-cross-encoder"
+
+
+class CodeBERTSQLErrorClassifier:
+    """CodeBERT cross-encoder inference wrapper."""
+
+    def __init__(
+        self,
+        model_dir: Union[str, Path] = DEFAULT_MODEL_DIR,
+        threshold: float = 0.5,
+        device: Optional[str] = None,
+    ):
+        self.model_dir = Path(model_dir)
+        self.threshold = threshold
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+        config_path = self.model_dir / "label_config.json"
+        if config_path.exists():
+            with open(config_path) as f:
+                cfg = json.load(f)
+            self.label_list = cfg.get("labels", load_codebert_labels())
+            self.threshold = cfg.get("threshold", threshold)
+            self.max_length = cfg.get("max_length", 512)
+        else:
+            self.label_list = load_codebert_labels()
+            self.max_length = 512
+
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            self.model_dir
+        ).to(self.device)
+        self.model.eval()
+
+    def predict(
+        self,
+        question: str,
+        schema: str,
+        student_sql: str,
+        correct_sql: str,
+        threshold: Optional[float] = None,
+        top_k: int = 5,
+    ) -> dict:
+        text = format_cross_encoder_input(
+            question=question,
+            schema=schema,
+            student_sql=student_sql,
+            correct_sql=correct_sql,
+        )
+        encoded = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_length,
+            padding=True,
+            return_tensors="pt",
+        ).to(self.device)
+
+        with torch.no_grad():
+            logits = self.model(**encoded).logits.cpu().numpy()[0]
+
+        probs = sigmoid(logits)
+        thr = threshold if threshold is not None else self.threshold
+        predicted = multihot_to_label_names(probs, self.label_list, threshold=thr)
+
+        ranked = sorted(
+            zip(self.label_list, probs.tolist()),
+            key=lambda x: x[1],
+            reverse=True,
+        )[:top_k]
+
+        return {
+            "error_labels": predicted,
+            "probabilities": {name: float(p) for name, p in ranked},
+            "top_k": [
+                {"label": name, "probability": float(p)} for name, p in ranked
+            ],
+            "primary_label": ranked[0][0],
+            "primary_confidence": float(ranked[0][1]),
+        }
+
+    def predict_batch(
+        self,
+        examples: List[dict],
+        batch_size: int = 16,
+    ) -> List[dict]:
+        results = []
+        for i in range(0, len(examples), batch_size):
+            chunk = examples[i : i + batch_size]
+            texts = [
+                format_cross_encoder_input(
+                    question=x["question"],
+                    schema=x["schema"],
+                    student_sql=x["student_sql"],
+                    correct_sql=x["correct_sql"],
+                )
+                for x in chunk
+            ]
+            encoded = self.tokenizer(
+                texts,
+                truncation=True,
+                max_length=self.max_length,
+                padding=True,
+                return_tensors="pt",
+            ).to(self.device)
+
+            with torch.no_grad():
+                logits = self.model(**encoded).logits.cpu().numpy()
+
+            for j, row in enumerate(logits):
+                probs = sigmoid(row)
+                results.append(
+                    {
+                        "error_labels": multihot_to_label_names(
+                            probs, self.label_list, self.threshold
+                        ),
+                        "primary_label": self.label_list[int(np.argmax(probs))],
+                        "primary_confidence": float(np.max(probs)),
+                    }
+                )
+        return results
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="CodeBERT SQL error inference")
+    parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
+    parser.add_argument("--question", type=str, required=True)
+    parser.add_argument("--schema", type=str, required=True)
+    parser.add_argument("--student-sql", type=str, required=True)
+    parser.add_argument("--correct-sql", type=str, required=True)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    args = parser.parse_args()
+
+    clf = CodeBERTSQLErrorClassifier(args.model_dir, threshold=args.threshold)
+    result = clf.predict(
+        question=args.question,
+        schema=args.schema,
+        student_sql=args.student_sql,
+        correct_sql=args.correct_sql,
+    )
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

src/hf_train_codebert.py

@@ -0,0 +1,226 @@
+"""Train CodeBERT cross-encoder for SQL error classification with HF Trainer."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    EarlyStoppingCallback,
+    Trainer,
+    TrainingArguments,
+)
+
+from src.codebert_dataset import (
+    SQLCodeBERTDataCollator,
+    prepare_datasets,
+)
+from src.codebert_labels import load_codebert_labels
+from src.hf_metrics import build_compute_metrics, compute_multilabel_metrics
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_DATA = PROJECT_ROOT / "data" / "sql_errors_1m.parquet"
+DEFAULT_OUTPUT = PROJECT_ROOT / "models" / "codebert-cross-encoder"
+DEFAULT_MODEL = "microsoft/codebert-base"
+
+
+def train(
+    data_path: Path | None = DEFAULT_DATA,
+    dataframe: pd.DataFrame | None = None,
+    output_dir: Path = DEFAULT_OUTPUT,
+    model_name: str = DEFAULT_MODEL,
+    epochs: float = 3.0,
+    batch_size: int = 16,
+    eval_batch_size: int = 32,
+    learning_rate: float = 2e-5,
+    weight_decay: float = 0.01,
+    warmup_ratio: float = 0.06,
+    max_length: int = 512,
+    max_samples: int | None = None,
+    test_size: float = 0.1,
+    val_size: float = 0.1,
+    threshold: float = 0.5,
+    seed: int = 42,
+    push_to_hub: bool = False,
+    hub_model_id: str | None = None,
+    fp16: bool = False,
+    save_strategy: str = "no",
+    hub_token: str | None = None,
+) -> dict:
+    if dataframe is not None:
+        df = dataframe.copy()
+        print(f"Loaded dataframe with {len(df):,} rows")
+    elif data_path is not None:
+        print(f"Loading dataset from {data_path}...")
+        df = pd.read_parquet(data_path)
+    else:
+        raise ValueError("Either data_path or dataframe must be provided")
+    if max_samples and len(df) > max_samples:
+        df = df.sample(n=max_samples, random_state=seed)
+
+    label_list = load_codebert_labels()
+    num_labels = len(label_list)
+    print(f"Labels ({num_labels}): {label_list}")
+    print(f"Samples: {len(df):,}")
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name,
+        num_labels=num_labels,
+        problem_type="multi_label_classification",
+        id2label={i: name for i, name in enumerate(label_list)},
+        label2id={name: i for i, name in enumerate(label_list)},
+    )
+
+    train_ds, val_ds, test_ds = prepare_datasets(
+        df,
+        tokenizer,
+        test_size=test_size,
+        val_size=val_size,
+        max_length=max_length,
+        seed=seed,
+    )
+    print(f"Train: {len(train_ds):,} | Val: {len(val_ds):,} | Test: {len(test_ds):,}")
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    label_info = {
+        "labels": label_list,
+        "model_name": model_name,
+        "architecture": "codebert-cross-encoder",
+        "input_format": "QUESTION + SCHEMA + STUDENT_SQL + CORRECT_SQL",
+        "max_length": max_length,
+        "threshold": threshold,
+    }
+    with open(output_dir / "label_config.json", "w") as f:
+        json.dump(label_info, f, indent=2)
+
+    training_args = TrainingArguments(
+        output_dir=str(output_dir),
+        num_train_epochs=epochs,
+        per_device_train_batch_size=batch_size,
+        per_device_eval_batch_size=eval_batch_size,
+        learning_rate=learning_rate,
+        weight_decay=weight_decay,
+        warmup_ratio=warmup_ratio,
+        eval_strategy="epoch",
+        save_strategy=save_strategy,
+        logging_strategy="steps",
+        logging_steps=50,
+        load_best_model_at_end=save_strategy == "epoch",
+        metric_for_best_model="f1_macro",
+        greater_is_better=True,
+        save_total_limit=1,
+        seed=seed,
+        report_to="none",
+        fp16=fp16 and torch.cuda.is_available(),
+        push_to_hub=push_to_hub,
+        hub_model_id=hub_model_id,
+        hub_token=hub_token,
+    )
+
+    callbacks = []
+    if save_strategy == "epoch":
+        callbacks.append(EarlyStoppingCallback(early_stopping_patience=2))
+
+    trainer_kwargs = dict(
+        model=model,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=val_ds,
+        data_collator=SQLCodeBERTDataCollator(tokenizer),
+        compute_metrics=build_compute_metrics(threshold=threshold),
+        callbacks=callbacks,
+    )
+    try:
+        trainer = Trainer(processing_class=tokenizer, **trainer_kwargs)
+    except TypeError:
+        trainer = Trainer(tokenizer=tokenizer, **trainer_kwargs)
+
+    print("Starting CodeBERT cross-encoder training...")
+    train_result = trainer.train()
+
+    print("Evaluating on validation set...")
+    val_metrics = trainer.evaluate()
+
+    print("Evaluating on held-out test set...")
+    test_output = trainer.predict(test_ds)
+    test_metrics = compute_multilabel_metrics(
+        test_output.predictions,
+        test_output.label_ids,
+        threshold=threshold,
+    )
+
+    trainer.save_model(str(output_dir))
+    tokenizer.save_pretrained(str(output_dir))
+
+    metrics = {
+        "train_samples": len(train_ds),
+        "val_samples": len(val_ds),
+        "test_samples": len(test_ds),
+        "train_runtime": train_result.metrics.get("train_runtime"),
+        "validation": val_metrics,
+        "test": test_metrics,
+    }
+    with open(output_dir / "metrics.json", "w") as f:
+        json.dump(metrics, f, indent=2, default=float)
+
+    print(f"\nValidation F1 (macro): {val_metrics.get('eval_f1_macro', 0):.4f}")
+    print(f"Test F1 (macro):       {test_metrics['f1_macro']:.4f}")
+    print(f"Test subset accuracy:  {test_metrics['subset_accuracy']:.4f}")
+    print(f"Model saved to {output_dir}")
+    return metrics
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Train CodeBERT cross-encoder with Hugging Face Trainer"
+    )
+    parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
+    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
+    parser.add_argument("--model-name", type=str, default=DEFAULT_MODEL)
+    parser.add_argument("--epochs", type=float, default=3.0)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--eval-batch-size", type=int, default=32)
+    parser.add_argument("--learning-rate", type=float, default=2e-5)
+    parser.add_argument("--max-length", type=int, default=512)
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument("--threshold", type=float, default=0.5)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--push-to-hub", action="store_true")
+    parser.add_argument("--hub-model-id", type=str, default=None)
+    parser.add_argument("--fp16", action="store_true")
+    parser.add_argument(
+        "--save-strategy",
+        choices=["no", "epoch"],
+        default="no",
+        help="Use 'no' to save only final model (saves disk space)",
+    )
+    args = parser.parse_args()
+
+    train(
+        data_path=args.data,
+        output_dir=args.output_dir,
+        model_name=args.model_name,
+        epochs=args.epochs,
+        batch_size=args.batch_size,
+        eval_batch_size=args.eval_batch_size,
+        learning_rate=args.learning_rate,
+        max_length=args.max_length,
+        max_samples=args.max_samples,
+        threshold=args.threshold,
+        seed=args.seed,
+        push_to_hub=args.push_to_hub,
+        hub_model_id=args.hub_model_id,
+        fp16=args.fp16,
+        save_strategy=args.save_strategy,
+    )
+
+
+if __name__ == "__main__":
+    main()

src/huggingface.py

@@ -0,0 +1,210 @@
+"""Hugging Face Hub integration for the SQL error classifier."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional, Union
+
+import joblib
+
+from src.categories import load_categories
+from src.cross_encoder_model import CrossEncoderClassifier
+from src.model import DEFAULT_ENCODER, load_model
+from src.multi_tower_model import MultiTowerClassifier, QueryContext
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+CONFIG_NAME = "config.json"
+CLASSIFIER_NAME = "classifier.joblib"
+CATEGORIES_NAME = "categories.json"
+
+SUPPORTED_CONTEXT_MODELS = (CrossEncoderClassifier, MultiTowerClassifier)
+
+
+class SQLLErrorClassifierHF:
+    """
+    Hugging Face–compatible wrapper for SQL error classifiers.
+
+    Usage:
+        clf = SQLLErrorClassifierHF.from_pretrained("username/sql-error-classifier")
+        result = clf.predict(
+            question="...", schema="...", correct_query="...", student_query="..."
+        )
+    """
+
+    def __init__(self, model, label_map: Dict[int, str]):
+        self.model = model
+        self.label_map = label_map
+
+    def predict(
+        self,
+        question: str,
+        schema: str,
+        correct_query: str,
+        student_query: str,
+        error_message: Optional[str] = None,
+        top_k: int = 3,
+    ) -> Dict[str, Any]:
+        ctx = QueryContext(
+            question=question,
+            schema=schema,
+            correct_query=correct_query,
+            student_query=student_query,
+            error_message=error_message,
+        )
+        proba = self.model.predict_proba([ctx])[0]
+        classes = self.model.classes_
+        ranked = sorted(zip(classes, proba), key=lambda x: x[1], reverse=True)
+        best_id = int(ranked[0][0])
+
+        diagnostics: Dict[str, Any] = {}
+        if isinstance(self.model, CrossEncoderClassifier):
+            diagnostics["pair_scores"] = self.model.explain_pair_scores(ctx)
+        elif isinstance(self.model, MultiTowerClassifier):
+            diagnostics["similarities"] = self.model.explain_similarities(ctx)
+
+        return {
+            "label_id": best_id,
+            "label_name": self.label_map[best_id],
+            "confidence": float(ranked[0][1]),
+            "top_k": [
+                {
+                    "label_id": int(cls),
+                    "label_name": self.label_map[int(cls)],
+                    "confidence": float(p),
+                }
+                for cls, p in ranked[:top_k]
+            ],
+            **diagnostics,
+        }
+
+    def save_pretrained(self, save_directory: Union[str, Path]) -> Path:
+        """Save model artifacts in Hugging Face Hub layout."""
+        save_dir = Path(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        if isinstance(self.model, CrossEncoderClassifier):
+            payload = {
+                "model_type": "cross_encoder",
+                "cross_encoder_name": self.model.cross_encoder_name,
+                "batch_size": self.model.batch_size,
+                "max_length": self.model.max_length,
+                "scaler": self.model.scaler,
+                "classifier": self.model.clf,
+                "classes_": self.model.classes_,
+            }
+            config = {
+                "model_type": "cross_encoder",
+                "architecture": "cross-encoder-pairwise",
+                "cross_encoder_name": self.model.cross_encoder_name,
+                "batch_size": self.model.batch_size,
+                "num_labels": len(self.label_map),
+                "task": "sql-error-classification",
+            }
+        elif isinstance(self.model, MultiTowerClassifier):
+            payload = {
+                "model_type": "multi_tower",
+                "encoder_name": self.model.encoder_name,
+                "batch_size": self.model.batch_size,
+                "scaler": self.model.scaler,
+                "classifier": self.model.clf,
+                "classes_": self.model.classes_,
+            }
+            config = {
+                "model_type": "multi_tower",
+                "architecture": "multi-tower-semantic-comparison",
+                "encoder_name": self.model.encoder_name,
+                "batch_size": self.model.batch_size,
+                "num_labels": len(self.label_map),
+                "task": "sql-error-classification",
+            }
+        else:
+            raise ValueError("Only cross_encoder and multi_tower models can be published")
+
+        joblib.dump(payload, save_dir / CLASSIFIER_NAME)
+        with open(save_dir / CONFIG_NAME, "w") as f:
+            json.dump(config, f, indent=2)
+
+        categories = load_categories()
+        cat_data = [
+            {"id": c.id, "name": c.name, "description": c.description}
+            for c in categories
+        ]
+        with open(save_dir / CATEGORIES_NAME, "w") as f:
+            json.dump(cat_data, f, indent=2)
+
+        return save_dir
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, Path],
+        *,
+        token: Optional[str] = None,
+    ) -> "SQLLErrorClassifierHF":
+        """Load from a local directory or Hugging Face Hub repo."""
+        path = _resolve_model_path(pretrained_model_name_or_path, token=token)
+
+        with open(path / CONFIG_NAME) as f:
+            config = json.load(f)
+
+        with open(path / CATEGORIES_NAME) as f:
+            categories = json.load(f)
+        label_map = {c["id"]: c["name"] for c in categories}
+
+        obj = joblib.load(path / CLASSIFIER_NAME)
+        model_type = config.get("model_type", obj.get("model_type"))
+
+        if model_type == "cross_encoder":
+            model = CrossEncoderClassifier(
+                cross_encoder_name=obj.get(
+                    "cross_encoder_name",
+                    config.get("cross_encoder_name", "cross-encoder/ms-marco-MiniLM-L6-v2"),
+                ),
+                batch_size=obj.get("batch_size", 32),
+                max_length=obj.get("max_length", 512),
+            )
+            model.scaler = obj["scaler"]
+            model.clf = obj["classifier"]
+            model.classes_ = obj.get("classes_", obj["classifier"].classes_)
+        else:
+            model = MultiTowerClassifier(
+                encoder_name=obj.get("encoder_name", config.get("encoder_name", DEFAULT_ENCODER)),
+                batch_size=obj.get("batch_size", 256),
+            )
+            model.scaler = obj["scaler"]
+            model.clf = obj["classifier"]
+            model.classes_ = obj.get("classes_", obj["classifier"].classes_)
+
+        return cls(model=model, label_map=label_map)
+
+
+def _resolve_model_path(
+    pretrained_model_name_or_path: Union[str, Path],
+    token: Optional[str] = None,
+) -> Path:
+    local = Path(pretrained_model_name_or_path)
+    if local.exists() and (local / CONFIG_NAME).exists():
+        return local
+
+    from huggingface_hub import snapshot_download
+
+    return Path(
+        snapshot_download(
+            repo_id=str(pretrained_model_name_or_path),
+            token=token,
+            allow_patterns=[CONFIG_NAME, CLASSIFIER_NAME, CATEGORIES_NAME],
+        )
+    )
+
+
+def package_for_hub(model_path: Path, output_dir: Path) -> Path:
+    """Convert a local joblib model into HF Hub layout."""
+    sklearn_model = load_model(model_path)
+    if not isinstance(sklearn_model, SUPPORTED_CONTEXT_MODELS):
+        raise ValueError(
+            "Only cross_encoder and multi_tower models can be published to Hugging Face Hub"
+        )
+    label_map = {c.id: c.name for c in load_categories()}
+    wrapper = SQLLErrorClassifierHF(model=sklearn_model, label_map=label_map)
+    return wrapper.save_pretrained(output_dir)

src/model.py

@@ -0,0 +1,321 @@
+"""SQL error classifiers: TF-IDF baseline and MiniLM embedding model."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import List, Literal, Optional, Protocol, Union
+
+import joblib
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.linear_model import SGDClassifier
+from sklearn.pipeline import FeatureUnion, Pipeline
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_MODEL_PATH = PROJECT_ROOT / "models" / "sql_error_classifier.joblib"
+DEFAULT_ENCODER = "sentence-transformers/all-MiniLM-L6-v2"
+
+ModelType = Literal["cross_encoder", "cross_encoder_ft", "multi_tower", "minilm", "tfidf"]
+
+
+class TextClassifier(Protocol):
+    classes_: np.ndarray
+
+    def fit(self, texts: List[str], y: np.ndarray) -> "TextClassifier": ...
+    def predict(self, texts: List[str]) -> np.ndarray: ...
+    def predict_proba(self, texts: List[str]) -> np.ndarray: ...
+
+
+def combine_features(
+    queries: List[str],
+    error_messages: Optional[List[str]] = None,
+    schemas: Optional[List[str]] = None,
+    questions: Optional[List[str]] = None,
+) -> List[str]:
+    """Fuse question, schema, query, and optional error message."""
+    texts: List[str] = []
+    for i, query in enumerate(queries):
+        parts: List[str] = []
+        if questions and questions[i]:
+            parts.append(f"QUESTION: {questions[i]}")
+        if schemas and schemas[i]:
+            parts.append(f"SCHEMA: {schemas[i]}")
+        parts.append(f"QUERY: {query}")
+        if error_messages and error_messages[i]:
+            parts.append(f"ERROR: {error_messages[i]}")
+        texts.append(" ".join(parts))
+    return texts
+
+
+def _build_text_features() -> FeatureUnion:
+    return FeatureUnion(
+        [
+            (
+                "word",
+                TfidfVectorizer(
+                    analyzer="word",
+                    ngram_range=(1, 2),
+                    max_features=30_000,
+                    sublinear_tf=True,
+                    strip_accents="unicode",
+                    token_pattern=r"(?u)\b\w+\b|(?<=[=<>!])\S+",
+                ),
+            ),
+            (
+                "char",
+                TfidfVectorizer(
+                    analyzer="char_wb",
+                    ngram_range=(2, 5),
+                    max_features=20_000,
+                    sublinear_tf=True,
+                ),
+            ),
+        ]
+    )
+
+
+def build_tfidf_classifier() -> Pipeline:
+    """Bag-of-words baseline. Fast but no deep semantic understanding."""
+    clf = SGDClassifier(
+        loss="log_loss",
+        penalty="l2",
+        alpha=1e-5,
+        max_iter=1000,
+        tol=1e-3,
+        class_weight="balanced",
+        random_state=42,
+    )
+    return Pipeline([("tfidf", _build_text_features()), ("clf", clf)])
+
+
+class EmbeddingClassifier:
+    """
+    MiniLM sentence embeddings + linear classifier.
+
+    Understands question intent (e.g. 'average' vs wrong aggregate) because
+    the encoder models full sentence context, not isolated word counts.
+    """
+
+    def __init__(
+        self,
+        encoder_name: str = DEFAULT_ENCODER,
+        batch_size: int = 256,
+    ):
+        self.encoder_name = encoder_name
+        self.batch_size = batch_size
+        self.encoder = None
+        self.clf = SGDClassifier(
+            loss="log_loss",
+            penalty="l2",
+            alpha=1e-4,
+            max_iter=1000,
+            tol=1e-3,
+            class_weight="balanced",
+            random_state=42,
+        )
+        self.classes_: Optional[np.ndarray] = None
+
+    def _load_encoder(self):
+        if self.encoder is None:
+            from sentence_transformers import SentenceTransformer
+
+            self.encoder = SentenceTransformer(self.encoder_name)
+
+    def encode(self, texts: List[str], show_progress: bool = False) -> np.ndarray:
+        self._load_encoder()
+        return self.encoder.encode(
+            texts,
+            batch_size=self.batch_size,
+            show_progress_bar=show_progress,
+            convert_to_numpy=True,
+        )
+
+    def fit(self, texts: List[str], y: np.ndarray) -> "EmbeddingClassifier":
+        X = self.encode(texts, show_progress=True)
+        self.clf.fit(X, y)
+        self.classes_ = self.clf.classes_
+        return self
+
+    def predict(self, texts: List[str]) -> np.ndarray:
+        return self.clf.predict(self.encode(texts))
+
+    def predict_proba(self, texts: List[str]) -> np.ndarray:
+        return self.clf.predict_proba(self.encode(texts))
+
+
+def build_classifier(
+    model_type: ModelType = "cross_encoder",
+) -> Union[
+    Pipeline,
+    EmbeddingClassifier,
+    "MultiTowerClassifier",
+    "CrossEncoderClassifier",
+    "FineTunedCrossEncoderClassifier",
+]:
+    if model_type == "tfidf":
+        return build_tfidf_classifier()
+    if model_type == "minilm":
+        return EmbeddingClassifier()
+    if model_type == "multi_tower":
+        from src.multi_tower_model import MultiTowerClassifier
+
+        return MultiTowerClassifier()
+    if model_type == "cross_encoder":
+        from src.cross_encoder_model import CrossEncoderClassifier
+
+        return CrossEncoderClassifier()
+    if model_type == "cross_encoder_ft":
+        from src.cross_encoder_model import FineTunedCrossEncoderClassifier
+
+        return FineTunedCrossEncoderClassifier()
+    raise ValueError(f"Unknown model_type: {model_type}")
+
+
+def save_model(
+    model: Union[
+        Pipeline,
+        EmbeddingClassifier,
+        "MultiTowerClassifier",
+        "CrossEncoderClassifier",
+        "FineTunedCrossEncoderClassifier",
+    ],
+    path: Path = DEFAULT_MODEL_PATH,
+    model_type: ModelType = "cross_encoder",
+) -> Path:
+    from src.cross_encoder_model import (
+        CrossEncoderClassifier,
+        FineTunedCrossEncoderClassifier,
+    )
+    from src.multi_tower_model import MultiTowerClassifier
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    if isinstance(model, FineTunedCrossEncoderClassifier):
+        ft_path = path if path.is_dir() or str(path).endswith("/") else path.with_suffix(".ce")
+        if ft_path.suffix == ".joblib":
+            ft_path = ft_path.with_suffix(".ce")
+        model.save(ft_path)
+        meta_path = ft_path / "meta.json" if ft_path.is_dir() else path.with_suffix(".meta.json")
+        with open(meta_path, "w") as f:
+            json.dump({"model_type": "cross_encoder_ft", "path": str(ft_path)}, f, indent=2)
+        return ft_path
+    if isinstance(model, CrossEncoderClassifier):
+        payload = {
+            "model_type": "cross_encoder",
+            "cross_encoder_name": model.cross_encoder_name,
+            "batch_size": model.batch_size,
+            "max_length": model.max_length,
+            "scaler": model.scaler,
+            "classifier": model.clf,
+            "classes_": model.classes_,
+        }
+        joblib.dump(payload, path)
+        meta_path = path.with_suffix(".meta.json")
+        with open(meta_path, "w") as f:
+            json.dump(
+                {
+                    "model_type": "cross_encoder",
+                    "cross_encoder_name": model.cross_encoder_name,
+                },
+                f,
+                indent=2,
+            )
+    elif isinstance(model, MultiTowerClassifier):
+        payload = {
+            "model_type": "multi_tower",
+            "encoder_name": model.encoder_name,
+            "batch_size": model.batch_size,
+            "scaler": model.scaler,
+            "classifier": model.clf,
+            "classes_": model.classes_,
+        }
+        joblib.dump(payload, path)
+        meta_path = path.with_suffix(".meta.json")
+        with open(meta_path, "w") as f:
+            json.dump(
+                {"model_type": "multi_tower", "encoder_name": model.encoder_name},
+                f,
+                indent=2,
+            )
+    elif isinstance(model, EmbeddingClassifier):
+        payload = {
+            "model_type": model_type,
+            "encoder_name": model.encoder_name,
+            "batch_size": model.batch_size,
+            "classifier": model.clf,
+            "classes_": model.classes_,
+        }
+        joblib.dump(payload, path)
+        meta_path = path.with_suffix(".meta.json")
+        with open(meta_path, "w") as f:
+            json.dump(
+                {"model_type": model_type, "encoder_name": model.encoder_name},
+                f,
+                indent=2,
+            )
+    else:
+        joblib.dump({"model_type": "tfidf", "pipeline": model}, path)
+    return path
+
+
+def load_model(
+    path: Path = DEFAULT_MODEL_PATH,
+) -> Union[
+    Pipeline,
+    EmbeddingClassifier,
+    "MultiTowerClassifier",
+    "CrossEncoderClassifier",
+    "FineTunedCrossEncoderClassifier",
+]:
+    from src.cross_encoder_model import (
+        CrossEncoderClassifier,
+        FineTunedCrossEncoderClassifier,
+    )
+    from src.multi_tower_model import MultiTowerClassifier
+
+    path = Path(path)
+
+    # Fine-tuned cross-encoder saved as directory
+    ce_path = path.with_suffix(".ce") if path.suffix == ".joblib" else path
+    if ce_path.exists() and (ce_path / "config.json").exists():
+        return FineTunedCrossEncoderClassifier.load(ce_path)
+
+    meta_path = path.with_suffix(".meta.json")
+    if meta_path.exists():
+        with open(meta_path) as f:
+            meta = json.load(f)
+        if meta.get("model_type") == "cross_encoder_ft":
+            ft_path = Path(meta.get("path", str(ce_path)))
+            return FineTunedCrossEncoderClassifier.load(ft_path)
+
+    obj = joblib.load(path)
+    if isinstance(obj, dict):
+        if obj.get("model_type") == "cross_encoder":
+            model = CrossEncoderClassifier(
+                cross_encoder_name=obj["cross_encoder_name"],
+                batch_size=obj.get("batch_size", 32),
+                max_length=obj.get("max_length", 512),
+            )
+            model.scaler = obj["scaler"]
+            model.clf = obj["classifier"]
+            model.classes_ = obj.get("classes_", obj["classifier"].classes_)
+            return model
+        if obj.get("model_type") == "multi_tower":
+            model = MultiTowerClassifier(
+                encoder_name=obj["encoder_name"],
+                batch_size=obj.get("batch_size", 256),
+            )
+            model.scaler = obj["scaler"]
+            model.clf = obj["classifier"]
+            model.classes_ = obj.get("classes_", obj["classifier"].classes_)
+            return model
+        if obj.get("model_type") == "minilm":
+            model = EmbeddingClassifier(
+                encoder_name=obj["encoder_name"],
+                batch_size=obj.get("batch_size", 256),
+            )
+            model.clf = obj["classifier"]
+            model.classes_ = obj.get("classes_", obj["classifier"].classes_)
+            return model
+        return obj["pipeline"]
+    return obj

src/multi_tower_model.py

@@ -0,0 +1,175 @@
+"""Multi-tower semantic comparison architecture for SQL error classification."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import numpy as np
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import StandardScaler
+
+from src.model import DEFAULT_ENCODER
+from src.sql_features import extract_sql_features
+
+
+@dataclass
+class QueryContext:
+    """Inputs available in the SQL playground at inference time."""
+
+    question: str
+    schema: str
+    correct_query: str
+    student_query: str
+    error_message: Optional[str] = None
+
+
+def _cosine(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    denom = np.linalg.norm(a, axis=1) * np.linalg.norm(b, axis=1)
+    denom = np.maximum(denom, 1e-8)
+    return np.sum(a * b, axis=1) / denom
+
+
+class MultiTowerClassifier:
+    """
+    Recommended architecture for SQL error classification.
+
+    Three semantic towers (shared MiniLM encoder):
+      1. Intent tower   — question + schema  → what should be answered
+      2. Reference tower — correct_query     → ground-truth solution
+      3. Student tower  — student_query      → what the student wrote
+
+    Comparison layer fuses:
+      - tower embeddings
+      - |student − reference|  (what changed)
+      - student ⊙ reference    (interaction)
+      - cosine similarities    (semantic alignment)
+      - SQL structural features (join/null/agg rules)
+
+    A light linear head maps the fused vector → 15 error categories.
+    """
+
+    def __init__(
+        self,
+        encoder_name: str = DEFAULT_ENCODER,
+        batch_size: int = 256,
+    ):
+        self.encoder_name = encoder_name
+        self.batch_size = batch_size
+        self.encoder = None
+        self.scaler = StandardScaler()
+        self.clf = LogisticRegression(
+            max_iter=1000,
+            solver="lbfgs",
+            class_weight="balanced",
+            random_state=42,
+        )
+        self.classes_: Optional[np.ndarray] = None
+
+    def _load_encoder(self):
+        if self.encoder is None:
+            from sentence_transformers import SentenceTransformer
+
+            self.encoder = SentenceTransformer(self.encoder_name)
+
+    def _encode(self, texts: List[str], show_progress: bool = False) -> np.ndarray:
+        self._load_encoder()
+        return self.encoder.encode(
+            texts,
+            batch_size=self.batch_size,
+            show_progress_bar=show_progress,
+            convert_to_numpy=True,
+        )
+
+    @staticmethod
+    def _intent_text(ctx: QueryContext) -> str:
+        return f"QUESTION: {ctx.question} SCHEMA: {ctx.schema}"
+
+    @staticmethod
+    def _reference_text(ctx: QueryContext) -> str:
+        return f"REFERENCE: {ctx.correct_query}"
+
+    @staticmethod
+    def _student_text(ctx: QueryContext) -> str:
+        parts = [f"STUDENT: {ctx.student_query}"]
+        if ctx.error_message:
+            parts.append(f"ERROR: {ctx.error_message}")
+        return " ".join(parts)
+
+    def _build_feature_matrix(
+        self,
+        contexts: List[QueryContext],
+        show_progress: bool = False,
+    ) -> np.ndarray:
+        intent_texts = [self._intent_text(c) for c in contexts]
+        ref_texts = [self._reference_text(c) for c in contexts]
+        student_texts = [self._student_text(c) for c in contexts]
+
+        intent_emb = self._encode(intent_texts, show_progress)
+        ref_emb = self._encode(ref_texts, show_progress=False)
+        student_emb = self._encode(student_texts, show_progress=False)
+
+        diff = np.abs(student_emb - ref_emb)
+        prod = student_emb * ref_emb
+        cos_sr = _cosine(student_emb, ref_emb).reshape(-1, 1)
+        cos_si = _cosine(student_emb, intent_emb).reshape(-1, 1)
+        cos_ri = _cosine(ref_emb, intent_emb).reshape(-1, 1)
+
+        sql_feats = np.array(
+            [
+                extract_sql_features(c.student_query, c.correct_query)
+                for c in contexts
+            ],
+            dtype=np.float64,
+        )
+
+        return np.hstack(
+            [intent_emb, ref_emb, student_emb, diff, prod, cos_sr, cos_si, cos_ri, sql_feats]
+        )
+
+    def fit(self, contexts: List[QueryContext], y: np.ndarray) -> "MultiTowerClassifier":
+        X = self._build_feature_matrix(contexts, show_progress=True)
+        X = self.scaler.fit_transform(X)
+        self.clf.fit(X, y)
+        self.classes_ = self.clf.classes_
+        return self
+
+    def _prepare_features(self, contexts: List[QueryContext]) -> np.ndarray:
+        X = self.scaler.transform(self._build_feature_matrix(contexts))
+        return np.nan_to_num(X, nan=0.0, posinf=1e3, neginf=-1e3)
+
+    def predict(self, contexts: List[QueryContext]) -> np.ndarray:
+        return self.clf.predict(self._prepare_features(contexts))
+
+    def predict_proba(self, contexts: List[QueryContext]) -> np.ndarray:
+        return self.clf.predict_proba(self._prepare_features(contexts))
+
+    def explain_similarities(self, ctx: QueryContext) -> dict:
+        """Diagnostic scores for the playground UI."""
+        emb = self._build_feature_matrix([ctx])
+        intent_texts = [self._intent_text(ctx)]
+        ref_texts = [self._reference_text(ctx)]
+        student_texts = [self._student_text(ctx)]
+        intent_emb = self._encode(intent_texts)
+        ref_emb = self._encode(ref_texts)
+        student_emb = self._encode(student_texts)
+        return {
+            "student_vs_reference": float(_cosine(student_emb, ref_emb)[0]),
+            "student_vs_intent": float(_cosine(student_emb, intent_emb)[0]),
+            "reference_vs_intent": float(_cosine(ref_emb, intent_emb)[0]),
+        }
+
+
+def contexts_from_dataframe(df) -> List[QueryContext]:
+    """Build QueryContext list from a training dataframe."""
+    has_error = "error_message" in df.columns
+    return [
+        QueryContext(
+            question=row["question"],
+            schema=row["schema"],
+            correct_query=row["correct_query"],
+            student_query=row["query"],
+            error_message=row["error_message"] if has_error else None,
+        )
+        for row in df.to_dict("records")
+    ]

src/predict.py

@@ -0,0 +1,132 @@
+"""Inference API for SQL error classification."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import List, Optional
+
+from src.categories import id_to_name, load_categories
+from src.model import DEFAULT_MODEL_PATH, combine_features, load_model
+from src.cross_encoder_model import (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+)
+from src.multi_tower_model import MultiTowerClassifier, QueryContext
+
+CONTEXT_MODELS = (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+    MultiTowerClassifier,
+)
+
+
+@dataclass
+class Prediction:
+    label_id: int
+    label_name: str
+    confidence: float
+    top_k: List[dict]
+    similarities: Optional[dict] = None
+    pair_scores: Optional[dict] = None
+
+
+class SQLErrorClassifier:
+    """Classifier wrapper for playground integration."""
+
+    def __init__(self, model_path: Path = DEFAULT_MODEL_PATH):
+        self.model = load_model(model_path)
+        self.label_map = id_to_name(load_categories())
+
+    def predict(
+        self,
+        query: str,
+        error_message: Optional[str] = None,
+        schema: Optional[str] = None,
+        question: Optional[str] = None,
+        correct_query: Optional[str] = None,
+        top_k: int = 3,
+    ) -> Prediction:
+        if isinstance(self.model, CONTEXT_MODELS):
+            if not all([schema, question, correct_query]):
+                raise ValueError(
+                    "context models require schema, question, and correct_query"
+                )
+            ctx = QueryContext(
+                question=question,
+                schema=schema,
+                correct_query=correct_query,
+                student_query=query,
+                error_message=error_message,
+            )
+            proba = self.model.predict_proba([ctx])[0]
+            similarities = (
+                self.model.explain_similarities(ctx)
+                if isinstance(self.model, MultiTowerClassifier)
+                else None
+            )
+            pair_scores = (
+                self.model.explain_pair_scores(ctx)
+                if isinstance(self.model, CrossEncoderClassifier)
+                else None
+            )
+        else:
+            pair_scores = None
+            similarities = None
+            text = combine_features(
+                queries=[query],
+                error_messages=[error_message] if error_message else None,
+                schemas=[schema] if schema else None,
+                questions=[question] if question else None,
+            )[0]
+            proba = self.model.predict_proba([text])[0]
+            similarities = None
+
+        classes = self.model.classes_
+        ranked = sorted(zip(classes, proba), key=lambda x: x[1], reverse=True)
+        best_id = int(ranked[0][0])
+
+        return Prediction(
+            label_id=best_id,
+            label_name=self.label_map[best_id],
+            confidence=float(ranked[0][1]),
+            top_k=[
+                {
+                    "label_id": int(cls),
+                    "label_name": self.label_map[int(cls)],
+                    "confidence": float(p),
+                }
+                for cls, p in ranked[:top_k]
+            ],
+            similarities=similarities,
+            pair_scores=pair_scores,
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Classify SQL error type")
+    parser.add_argument("--query", type=str, required=True)
+    parser.add_argument("--correct-query", type=str, default=None)
+    parser.add_argument("--error-message", type=str, default=None)
+    parser.add_argument("--schema", type=str, default=None)
+    parser.add_argument("--question", type=str, default=None)
+    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL_PATH)
+    parser.add_argument("--top-k", type=int, default=3)
+    args = parser.parse_args()
+
+    clf = SQLErrorClassifier(args.model)
+    result = clf.predict(
+        args.query,
+        args.error_message,
+        args.schema,
+        args.question,
+        args.correct_query,
+        top_k=args.top_k,
+    )
+    print(json.dumps(asdict(result), indent=2))
+
+
+if __name__ == "__main__":
+    main()

src/sql_features.py

@@ -0,0 +1,81 @@
+"""Lightweight structural SQL features for the classification head."""
+
+from __future__ import annotations
+
+import re
+from typing import List
+
+AGG_FUNCS = ("COUNT", "SUM", "AVG", "MAX", "MIN")
+WINDOW_FUNCS = ("ROW_NUMBER", "RANK", "DENSE_RANK", "OVER")
+
+
+def _upper(sql: str) -> str:
+    return sql.upper()
+
+
+def extract_sql_features(student_query: str, correct_query: str = "") -> List[float]:
+    """
+    Rule-based signals that complement semantic embeddings.
+    Returns a fixed-length float vector.
+    """
+    s = _upper(student_query)
+    c = _upper(correct_query) if correct_query else ""
+
+    has_agg = any(f" {f}(" in s or f"{f}(" in s for f in AGG_FUNCS)
+    has_group = "GROUP BY" in s
+    has_join = "JOIN" in s
+    has_on = " ON " in s
+    has_where = " WHERE " in s
+    has_having = " HAVING " in s
+    has_distinct = "DISTINCT" in s
+    has_subquery = "(" in s and "SELECT" in s[s.find("(") :]
+    has_window = "OVER" in s
+    has_null_eq = "= NULL" in s or "=NULL" in s
+    has_is_null = "IS NULL" in s or "IS NOT NULL" in s
+    has_select_star = bool(re.search(r"SELECT\s+\*", s))
+    has_or = " OR " in s
+    has_and = " AND " in s
+
+    correct_has_distinct = "DISTINCT" in c
+    correct_has_group = "GROUP BY" in c
+    correct_has_inner = "INNER JOIN" in c
+    student_has_left = "LEFT JOIN" in s
+
+    return [
+        float(has_agg),
+        float(has_agg and not has_group),
+        float(has_join and not has_on),
+        float(has_join),
+        float(has_where and has_having),
+        float(has_agg and has_where and not has_having),
+        float(has_distinct),
+        float(correct_has_distinct and not has_distinct),
+        float(has_subquery),
+        float(has_window),
+        float(has_null_eq),
+        float(has_is_null),
+        float(has_select_star),
+        float(has_or and has_and),
+        float(correct_has_inner and student_has_left),
+        float(len(s) / max(len(c), 1)),  # length ratio vs reference
+    ]
+
+
+FEATURE_NAMES = [
+    "has_aggregate",
+    "agg_without_group_by",
+    "join_without_on",
+    "has_join",
+    "where_and_having",
+    "agg_in_where",
+    "has_distinct",
+    "missing_distinct_vs_correct",
+    "has_subquery",
+    "has_window",
+    "null_equals",
+    "is_null_check",
+    "select_star",
+    "and_or_mix",
+    "left_vs_inner_join",
+    "length_ratio",
+]

src/sql_templates.py

@@ -0,0 +1,258 @@
+"""Error injectors that transform exercise context into labeled mistakes."""
+
+from __future__ import annotations
+
+import random
+from typing import Callable, Dict, List, Tuple
+
+from src.exercises import Exercise
+
+FAKE_COLUMNS = ["fullname", "studentname", "coursename", "dept_name", "totals"]
+FAKE_TABLES = ["student", "course", "enrolment", "employe", "orderz"]
+
+
+def _pick(rng: random.Random, items: List[str], k: int = 1) -> str | List[str]:
+    if k == 1:
+        return rng.choice(items)
+    return rng.sample(items, k)
+
+
+def _first_table(exercise: Exercise) -> str:
+    return exercise.tables[0]
+
+
+def _second_table(exercise: Exercise) -> str:
+    return exercise.tables[1] if len(exercise.tables) > 1 else exercise.tables[0]
+
+
+# --- Error injectors: (exercise) -> (erroneous_sql, error_message) ---
+
+def inject_syntax_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    sql = exercise.correct_query
+    mutations = [
+        lambda s: s.replace("SELECT", "SELEC", 1),
+        lambda s: s.replace("FROM", "FRO", 1),
+        lambda s: s[:-1],
+        lambda s: s.replace(")", "", 1),
+        lambda s: s + " WHERE",
+        lambda s: s.replace(",", "", 1),
+        lambda s: s.replace("'", '"', 1) if "'" in s else s + " 'unclosed",
+    ]
+    return rng.choice(mutations)(sql), "syntax error at or near unexpected token"
+
+
+def inject_join_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t1, t2 = _first_table(exercise), _second_table(exercise)
+    col = _pick(rng, list(exercise.columns))
+    variants = [
+        f"SELECT {col} FROM {t1} JOIN {t2}",
+        f"SELECT {col} FROM {t1} INNER JOIN {t2} ON {t1}.id = {t2}.id",
+        (
+            f"SELECT {t1}.{col} FROM {t1} "
+            f"LEFT JOIN {t2} ON {t1}.{col} = {t2}.{col}"
+        ),
+        f"SELECT * FROM {t1}, {t2} WHERE {t1}.wrong_id = {t2}.wrong_id",
+    ]
+    return rng.choice(variants), "missing ON clause or invalid join condition"
+
+
+def inject_aggregation_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    cols = list(exercise.columns)
+    group_col = cols[0]
+    agg_col = cols[-1]
+    bad = f"SELECT {group_col}, AVG({agg_col}) FROM {t}"
+    return bad, "column must appear in GROUP BY clause or be used in aggregate function"
+
+
+def inject_having_where_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    cols = list(exercise.columns)
+    group_col, agg_col = cols[0], cols[-1]
+    bad = (
+        f"SELECT {group_col}, COUNT({agg_col}) FROM {t} "
+        f"WHERE COUNT({agg_col}) > {rng.randint(1, 5)}"
+    )
+    return bad, "aggregate functions are not allowed in WHERE"
+
+
+def inject_subquery_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t1, t2 = _first_table(exercise), _second_table(exercise)
+    col = _pick(rng, list(exercise.columns))
+    bad = f"SELECT {col} FROM {t1} WHERE {col} = (SELECT {col} FROM {t2})"
+    return bad, "subquery returned more than one row"
+
+
+def inject_window_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    col = _pick(rng, list(exercise.columns))
+    variants = [
+        f"SELECT {col}, ROW_NUMBER() OVER () FROM {t}",
+        f"SELECT {col}, SUM({col}) OVER (ORDER BY {col}) FROM {t} GROUP BY {col}",
+        f"SELECT {col}, RANK() OVER (PARTITION {col}) FROM {t}",
+    ]
+    return rng.choice(variants), "window function requires PARTITION BY or ORDER BY"
+
+
+def inject_null_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    col = _pick(rng, list(exercise.columns))
+    bad = f"SELECT * FROM {t} WHERE {col} = NULL"
+    return bad, "use IS NULL or IS NOT NULL to test for null values"
+
+
+def inject_date_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    variants = [
+        f"SELECT * FROM {t} WHERE order_date = '31/02/2023'",
+        f"SELECT * FROM {t} WHERE order_date = DATE '2023-13-40'",
+        f"SELECT * FROM {t} WHERE STR_TO_DATE('bad-date', '%Y-%m-%d')",
+        f"SELECT * FROM {t} WHERE hire_date > 'yesterday'",
+    ]
+    return rng.choice(variants), "invalid date format or unknown date function"
+
+
+def inject_column_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    col = _pick(rng, FAKE_COLUMNS)
+    bad = f"SELECT {col} FROM {t}"
+    return bad, f"column '{col}' does not exist"
+
+
+def inject_table_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    tbl = _pick(rng, FAKE_TABLES)
+    col = _pick(rng, list(exercise.columns))
+    bad = f"SELECT {col} FROM {tbl}"
+    return bad, f"relation '{tbl}' does not exist"
+
+
+def inject_datatype_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t = _first_table(exercise)
+    col = _pick(rng, list(exercise.columns))
+    bad = f"SELECT {col} FROM {t} WHERE {col} = '{rng.choice(['abc', 'ten', 'N/A'])}'"
+    return bad, "operator does not exist: integer = character varying"
+
+
+def inject_duplicate_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    """Drop DISTINCT when the question asks for unique values."""
+    sql = exercise.correct_query
+    if "DISTINCT" in sql.upper():
+        bad = sql.upper().replace("DISTINCT ", "").replace("distinct ", "")
+        # restore original casing loosely
+        bad = sql.replace("DISTINCT ", "").replace("distinct ", "")
+    else:
+        col = _pick(rng, list(exercise.columns))
+        bad = f"SELECT {col} FROM {_first_table(exercise)}"
+    return bad, "query returns duplicate rows; DISTINCT may be required"
+
+
+def inject_logical_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    """
+    Produce a query that runs against the schema but answers the question incorrectly.
+    Variants are tied to the exercise question and correct answer.
+    """
+    sql = exercise.correct_query
+    q = exercise.question.lower()
+    variants: List[str] = []
+
+    if "average" in q or "avg" in sql.lower():
+        variants.append(sql.replace("AVG(", "SUM(", 1))
+        variants.append(sql.replace("AVG(", "MAX(", 1))
+    if " and " in q and " AND " in sql:
+        variants.append(sql.replace(" AND ", " OR ", 1))
+    if "join" in sql.lower():
+        t1, t2 = _first_table(exercise), _second_table(exercise)
+        variants.append(
+            f"SELECT {t1}.name, {t2}.name FROM {t1} "
+            f"JOIN {t2} ON {t1}.id = {t2}.id"
+        )
+        variants.append(sql.replace("INNER JOIN", "LEFT JOIN", 1))
+    if "between" in q and "BETWEEN" in sql.upper():
+        upper = sql.upper()
+        between_part = upper.split("BETWEEN", 1)[1]
+        bounds = between_part.split("AND", 1)
+        if len(bounds) == 2:
+            lo = bounds[0].strip().split()[-1]
+            hi = bounds[1].strip().split()[0]
+            variants.append(
+                sql.split("WHERE", 1)[0]
+                + f" WHERE price BETWEEN {hi} AND {lo}"
+            )
+    if "rank" in q or "over" in sql.lower():
+        col = _pick(rng, list(exercise.columns))
+        variants.append(
+            f"SELECT name, {col} FROM {_first_table(exercise)} ORDER BY {col} DESC"
+        )
+    if "total" in q and "WHERE" in sql.upper():
+        variants.append(sql.replace("active", "inactive"))
+    if "highest" in q or "max" in sql.lower():
+        col = _pick(rng, list(exercise.columns))
+        variants.append(
+            f"SELECT name FROM {_first_table(exercise)} "
+            f"WHERE {col} >= (SELECT AVG({col}) FROM {_second_table(exercise)})"
+        )
+    if "enrolled" in q and "INNER JOIN" in sql.upper():
+        variants.append(sql.replace("INNER JOIN", "LEFT JOIN", 1))
+    if "not provided" in q or "is null" in sql.lower():
+        variants.append(sql.replace("IS NULL", "= ''"))
+
+    if not variants:
+        col = _pick(rng, list(exercise.columns))
+        t = _first_table(exercise)
+        variants = [
+            f"SELECT {col} FROM {t} ORDER BY {col} DESC LIMIT 10",
+            f"SELECT COUNT(*) FROM {t}",
+        ]
+
+    bad = rng.choice(variants)
+    return bad, "query executes but produces incorrect result set"
+
+
+def inject_performance_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    t1, t2 = _first_table(exercise), _second_table(exercise)
+    variants = [
+        f"SELECT * FROM {t1}",
+        f"SELECT * FROM {t1} JOIN {t2} ON {t1}.id = {t2}.id",
+        (
+            f"SELECT * FROM {t1} "
+            f"WHERE {_pick(rng, list(exercise.columns))} "
+            f"LIKE '%{rng.choice(['a', 'e', 'i'])}%'"
+        ),
+        f"SELECT * FROM {t1} CROSS JOIN {t2}",
+    ]
+    return rng.choice(variants), "inefficient query: SELECT * or cartesian join detected"
+
+
+def inject_filtering_error(rng: random.Random, exercise: Exercise) -> Tuple[str, str]:
+    sql = exercise.correct_query
+    col = _pick(rng, list(exercise.columns))
+    t = _first_table(exercise)
+    threshold = rng.randint(50, 90)
+    variants = [
+        sql.replace(">", "<", 1) if ">" in sql else sql.replace("=", "!=", 1),
+        f"SELECT {col} FROM {t} WHERE {col} > {threshold} AND {col} < {threshold - 20}",
+        f"SELECT {col} FROM {t} WHERE NOT {col} > {threshold}",
+        sql.replace(" AND ", " OR ", 1) if " AND " in sql else (
+            f"SELECT {col} FROM {t} WHERE {col} BETWEEN {threshold} AND {threshold - 10}"
+        ),
+    ]
+    return rng.choice(variants), "WHERE clause filters incorrect rows"
+
+
+ERROR_INJECTORS: Dict[int, Callable[[random.Random, Exercise], Tuple[str, str]]] = {
+    0: inject_syntax_error,
+    1: inject_join_error,
+    2: inject_aggregation_error,
+    3: inject_having_where_error,
+    4: inject_subquery_error,
+    5: inject_window_error,
+    6: inject_null_error,
+    7: inject_date_error,
+    8: inject_column_error,
+    9: inject_table_error,
+    10: inject_datatype_error,
+    11: inject_duplicate_error,
+    12: inject_logical_error,
+    13: inject_performance_error,
+    14: inject_filtering_error,
+}

src/train.py

@@ -0,0 +1,198 @@
+"""Train the SQL error classifier."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import pandas as pd
+from sklearn.metrics import classification_report
+from sklearn.model_selection import train_test_split
+
+from src.categories import id_to_name, load_categories
+from src.cross_encoder_model import (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+)
+from src.model import (
+    DEFAULT_MODEL_PATH,
+    ModelType,
+    build_classifier,
+    combine_features,
+    save_model,
+)
+from src.multi_tower_model import MultiTowerClassifier, contexts_from_dataframe
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_DATA = PROJECT_ROOT / "data" / "sql_errors_1m.parquet"
+DEFAULT_METRICS = PROJECT_ROOT / "models" / "metrics.json"
+
+CONTEXT_MODELS = (
+    CrossEncoderClassifier,
+    FineTunedCrossEncoderClassifier,
+    MultiTowerClassifier,
+)
+
+
+def _split_dataframe(
+    df: pd.DataFrame, test_size: float, val_size: float, seed: int
+) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+    trainval, test = train_test_split(
+        df, test_size=test_size, random_state=seed, stratify=df["label_id"]
+    )
+    relative_val = val_size / (1 - test_size)
+    train, val = train_test_split(
+        trainval,
+        test_size=relative_val,
+        random_state=seed,
+        stratify=trainval["label_id"],
+    )
+    return train, val, test
+
+
+def train(
+    data_path: Path = DEFAULT_DATA,
+    model_path: Path = DEFAULT_MODEL_PATH,
+    metrics_path: Path = DEFAULT_METRICS,
+    test_size: float = 0.1,
+    val_size: float = 0.1,
+    use_error_message: bool = True,
+    max_train_samples: int | None = None,
+    model_type: ModelType = "cross_encoder",
+    epochs: int = 1,
+    seed: int = 42,
+) -> dict:
+    print(f"Loading data from {data_path}...")
+    df = pd.read_parquet(data_path)
+
+    if max_train_samples and len(df) > max_train_samples:
+        df = df.sample(n=max_train_samples, random_state=seed)
+
+    if not use_error_message and "error_message" in df.columns:
+        df = df.drop(columns=["error_message"])
+
+    train_df, val_df, test_df = _split_dataframe(df, test_size, val_size, seed)
+    print(
+        f"Train: {len(train_df):,} | Val: {len(val_df):,} | Test: {len(test_df):,}"
+    )
+
+    model = build_classifier(model_type=model_type)
+    print(f"Training {model_type} classifier...")
+
+    if isinstance(model, CONTEXT_MODELS):
+        train_ctx = contexts_from_dataframe(train_df)
+        val_ctx = contexts_from_dataframe(val_df)
+        test_ctx = contexts_from_dataframe(test_df)
+
+        if isinstance(model, FineTunedCrossEncoderClassifier):
+            model.fit(
+                train_ctx,
+                train_df["label_id"].values,
+                epochs=epochs,
+                output_path=model_path.with_suffix(".ce")
+                if model_path.suffix == ".joblib"
+                else model_path,
+            )
+        else:
+            model.fit(train_ctx, train_df["label_id"].values)
+
+        val_preds = model.predict(val_ctx)
+        test_preds = model.predict(test_ctx)
+        y_val = val_df["label_id"].values
+        y_test = test_df["label_id"].values
+    else:
+
+        def to_texts(frame: pd.DataFrame) -> list[str]:
+            return combine_features(
+                queries=frame["query"].tolist(),
+                error_messages=frame["error_message"].tolist()
+                if "error_message" in frame.columns
+                else None,
+                schemas=frame["schema"].tolist() if "schema" in frame.columns else None,
+                questions=frame["question"].tolist()
+                if "question" in frame.columns
+                else None,
+            )
+
+        model.fit(to_texts(train_df), train_df["label_id"].values)
+        val_preds = model.predict(to_texts(val_df))
+        test_preds = model.predict(to_texts(test_df))
+        y_val = val_df["label_id"].values
+        y_test = test_df["label_id"].values
+
+    val_report = classification_report(
+        y_val, val_preds, output_dict=True, zero_division=0
+    )
+    print(f"Validation accuracy: {val_report['accuracy']:.4f}")
+
+    test_report = classification_report(
+        y_test, test_preds, output_dict=True, zero_division=0
+    )
+    print(f"Test accuracy: {test_report['accuracy']:.4f}")
+
+    save_model(model, model_path, model_type=model_type)
+    print(f"Model saved to {model_path}")
+
+    categories = load_categories()
+    label_map = id_to_name(categories)
+    metrics = {
+        "train_size": len(train_df),
+        "val_size": len(val_df),
+        "test_size": len(test_df),
+        "model_type": model_type,
+        "epochs": epochs if model_type == "cross_encoder_ft" else None,
+        "use_error_message": use_error_message,
+        "validation": val_report,
+        "test": test_report,
+        "label_map": {str(k): v for k, v in label_map.items()},
+    }
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(metrics_path, "w") as f:
+        json.dump(metrics, f, indent=2)
+    print(f"Metrics saved to {metrics_path}")
+
+    return metrics
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train SQL error classifier")
+    parser.add_argument("--data", type=Path, default=DEFAULT_DATA)
+    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL_PATH)
+    parser.add_argument("--metrics", type=Path, default=DEFAULT_METRICS)
+    parser.add_argument("--test-size", type=float, default=0.1)
+    parser.add_argument("--val-size", type=float, default=0.1)
+    parser.add_argument("--no-error-message", action="store_true")
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument(
+        "--model-type",
+        choices=["cross_encoder", "cross_encoder_ft", "multi_tower", "minilm", "tfidf"],
+        default="cross_encoder",
+        help="cross_encoder (recommended): joint attention pairs; "
+        "cross_encoder_ft: fine-tuned end-to-end (best accuracy)",
+    )
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=1,
+        help="Epochs for cross_encoder_ft fine-tuning",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    train(
+        data_path=args.data,
+        model_path=args.model,
+        metrics_path=args.metrics,
+        test_size=args.test_size,
+        val_size=args.val_size,
+        use_error_message=not args.no_error_message,
+        max_train_samples=args.max_samples,
+        model_type=args.model_type,
+        epochs=args.epochs,
+        seed=args.seed,
+    )
+
+
+if __name__ == "__main__":
+    main()

train_space_app.py

@@ -0,0 +1,230 @@
+"""
+Hugging Face Space — CodeBERT SQL Error Classifier Training UI.
+
+Deploy as a Gradio Space with app_file: train_space_app.py
+Set hardware to GPU (t4-small recommended).
+Add HF_TOKEN secret to push trained models to your Hub account.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+import pandas as pd
+
+from src.hf_train_codebert import train
+
+PROJECT_ROOT = Path(__file__).parent
+DEFAULT_DATA = PROJECT_ROOT / "data" / "sql_errors_dev.parquet"
+OUTPUT_DIR = PROJECT_ROOT / "models" / "codebert-cross-encoder"
+BUNDLED_DATASETS = {
+    "Dev (15K samples)": str(PROJECT_ROOT / "data" / "sql_errors_dev.parquet"),
+    "Full (1M samples)": str(PROJECT_ROOT / "data" / "sql_errors_1m.parquet"),
+}
+
+
+def _format_metrics(metrics: dict) -> str:
+    val = metrics.get("validation", {})
+    test = metrics.get("test", {})
+    lines = [
+        "## Training complete",
+        "",
+        f"- Train samples: **{metrics.get('train_samples', 0):,}**",
+        f"- Val samples: **{metrics.get('val_samples', 0):,}**",
+        f"- Test samples: **{metrics.get('test_samples', 0):,}**",
+        "",
+        "### Validation",
+        f"- F1 macro: **{val.get('eval_f1_macro', 0):.4f}**",
+        f"- F1 micro: **{val.get('eval_f1_micro', 0):.4f}**",
+        "",
+        "### Test",
+        f"- F1 macro: **{test.get('f1_macro', 0):.4f}**",
+        f"- F1 micro: **{test.get('f1_micro', 0):.4f}**",
+        f"- Subset accuracy: **{test.get('subset_accuracy', 0):.4f}**",
+        "",
+        f"Model saved to `{OUTPUT_DIR}`",
+    ]
+    if metrics.get("hub_url"):
+        lines.append(f"\n**Hub model:** {metrics['hub_url']}")
+    return "\n".join(lines)
+
+
+def run_training(
+    dataset_choice: str,
+    uploaded_file,
+    max_samples: int,
+    epochs: float,
+    batch_size: int,
+    learning_rate: float,
+    max_length: int,
+    fp16: bool,
+    push_to_hub: bool,
+    hub_model_id: str,
+    progress=gr.Progress(),
+):
+    progress(0, desc="Preparing dataset...")
+
+    if uploaded_file is not None:
+        data_path = Path(uploaded_file.name)
+    else:
+        data_path = Path(BUNDLED_DATASETS.get(dataset_choice, DEFAULT_DATA))
+        if not data_path.exists():
+            return (
+                f"Dataset not found: `{data_path}`. "
+                "Upload a parquet file or include data/ in the Space repo.",
+                None,
+                None,
+            )
+
+    hub_token = os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
+    if push_to_hub and not hub_token:
+        return (
+            "Add `HF_TOKEN` to Space secrets to push models to the Hub.",
+            None,
+            None,
+        )
+    if push_to_hub and not hub_model_id.strip():
+        return "Enter a Hub model id (e.g. `your-username/sql-codebert-classifier`).", None, None
+
+    if OUTPUT_DIR.exists():
+        shutil.rmtree(OUTPUT_DIR, ignore_errors=True)
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    samples = int(max_samples) if max_samples and max_samples > 0 else None
+    progress(0.1, desc="Starting CodeBERT training...")
+
+    try:
+        metrics = train(
+            data_path=data_path,
+            output_dir=OUTPUT_DIR,
+            epochs=epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            max_length=max_length,
+            max_samples=samples,
+            fp16=fp16,
+            save_strategy="no",
+            push_to_hub=push_to_hub,
+            hub_model_id=hub_model_id.strip() or None,
+            hub_token=hub_token,
+        )
+    except Exception as exc:
+        return f"Training failed:\n\n```\n{exc}\n```", None, None
+
+    progress(1.0, desc="Done")
+    if push_to_hub and hub_model_id.strip():
+        metrics["hub_url"] = f"https://huggingface.co/{hub_model_id.strip()}"
+
+    metrics_path = OUTPUT_DIR / "metrics.json"
+    summary = _format_metrics(metrics)
+    return summary, str(metrics_path) if metrics_path.exists() else None, str(OUTPUT_DIR)
+
+
+def load_preview(dataset_choice: str, uploaded_file) -> str:
+    try:
+        if uploaded_file is not None:
+            df = pd.read_parquet(uploaded_file.name)
+        else:
+            path = BUNDLED_DATASETS.get(dataset_choice, DEFAULT_DATA)
+            if not Path(path).exists():
+                return f"Dataset not found: {path}"
+            df = pd.read_parquet(path)
+        cols = list(df.columns)
+        sample = df.head(2).to_dict(orient="records")
+        return f"**Rows:** {len(df):,}\n\n**Columns:** `{cols}`\n\n**Sample:**\n```json\n{json.dumps(sample, indent=2)[:2000]}\n```"
+    except Exception as exc:
+        return f"Could not load preview: {exc}"
+
+
+with gr.Blocks(title="SQL Error Classifier — Train") as demo:
+    gr.Markdown(
+        """
+        # SQL Error Classifier — CodeBERT Training
+        Train **microsoft/codebert-base** as a cross-encoder on this Space.
+
+        **Input format:** `QUESTION` + `SCHEMA` + `STUDENT_SQL` + `CORRECT_SQL` (single sequence)
+
+        **GPU recommended** — upgrade Space hardware to `t4-small` or better.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            dataset_choice = gr.Dropdown(
+                choices=list(BUNDLED_DATASETS.keys()),
+                value="Dev (15K samples)",
+                label="Bundled dataset",
+            )
+            uploaded = gr.File(
+                label="Or upload parquet",
+                file_types=[".parquet"],
+            )
+            preview_btn = gr.Button("Preview dataset")
+            preview_out = gr.Markdown()
+
+            max_samples = gr.Number(
+                label="Max samples (0 = all)",
+                value=5000,
+                precision=0,
+            )
+            epochs = gr.Slider(1, 10, value=2, step=1, label="Epochs")
+            batch_size = gr.Slider(4, 64, value=8, step=4, label="Batch size")
+            learning_rate = gr.Number(label="Learning rate", value=2e-5)
+            max_length = gr.Slider(128, 512, value=512, step=64, label="Max length")
+            fp16 = gr.Checkbox(label="FP16 (GPU only)", value=True)
+
+            push_to_hub = gr.Checkbox(label="Push to Hugging Face Hub", value=False)
+            hub_model_id = gr.Textbox(
+                label="Hub model id",
+                placeholder="your-username/sql-codebert-classifier",
+            )
+
+            train_btn = gr.Button("Start Training", variant="primary")
+
+        with gr.Column(scale=1):
+            result = gr.Markdown(label="Results")
+            metrics_file = gr.File(label="metrics.json")
+            model_dir = gr.Textbox(label="Model output path", interactive=False)
+
+    preview_btn.click(load_preview, [dataset_choice, uploaded], preview_out)
+    train_btn.click(
+        run_training,
+        [
+            dataset_choice,
+            uploaded,
+            max_samples,
+            epochs,
+            batch_size,
+            learning_rate,
+            max_length,
+            fp16,
+            push_to_hub,
+            hub_model_id,
+        ],
+        [result, metrics_file, model_dir],
+    )
+
+    gr.Markdown(
+        """
+        ### Space setup
+        1. Create a Gradio Space and push this repo
+        2. Set **Hardware → GPU (t4-small)**
+        3. Add secret `HF_TOKEN` (write token) to push models
+        4. Include `data/sql_errors_dev.parquet` in the repo (or upload at runtime)
+
+        ### After training
+        Use the saved model with:
+        ```python
+        from src.hf_predict_codebert import CodeBERTSQLErrorClassifier
+        clf = CodeBERTSQLErrorClassifier("models/codebert-cross-encoder")
+        ```
+        """
+    )
+
+if __name__ == "__main__":
+    demo.launch()