Deploy CodeBERT inference Space

.gitignore

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

Dockerfile

@@ -0,0 +1,17 @@
+# Inference Space — loads model from Hugging Face Hub at startup
+FROM python:3.12-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENV GRADIO_SERVER_NAME=0.0.0.0
+ENV GRADIO_SERVER_PORT=7860
+ENV SPACE_MODEL_ID=nishu08/sql-codebert-classifier
+
+EXPOSE 7860
+CMD ["python", "app.py"]

Dockerfile.inference

@@ -0,0 +1,17 @@
+# Inference Space — loads model from Hugging Face Hub at startup
+FROM python:3.12-slim
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+ENV GRADIO_SERVER_NAME=0.0.0.0
+ENV GRADIO_SERVER_PORT=7860
+ENV SPACE_MODEL_ID=nishu08/sql-codebert-classifier
+
+EXPOSE 7860
+CMD ["python", "app.py"]

README.md

@@ -1,10 +1,32 @@
 ---
-title: Sql Error Classifier
-emoji: ⚡
-colorFrom: indigo
-colorTo: pink
+title: SQL Error Classifier
+emoji: 🗄️
+colorFrom: blue
+colorTo: indigo
 sdk: docker
+app_port: 7860
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SQL Error Classifier — Inference
+
+Classify student SQL mistakes using a **CodeBERT cross-encoder** model.
+
+Model: [nishu08/sql-codebert-classifier](https://huggingface.co/nishu08/sql-codebert-classifier)
+
+## Inputs
+
+- **Question** — what the student was asked
+- **Schema** — available tables/columns
+- **Student SQL** — student's submission
+- **Correct SQL** — reference answer
+
+## Labels
+
+`JOIN_ERROR`, `AGGREGATION_ERROR`, `FILTER_ERROR`, `WINDOW_FUNCTION_ERROR`,
+`SUBQUERY_ERROR`, `NULL_HANDLING_ERROR`, `PERFORMANCE_ERROR`, `LOGICAL_ERROR`, `SYNTAX_ERROR`
+
+## Secrets (optional)
+
+- `SPACE_MODEL_ID` — override model repo (default: `nishu08/sql-codebert-classifier`)

README_INFERENCE_SPACE.md

@@ -0,0 +1,32 @@
+---
+title: SQL Error Classifier
+emoji: 🗄️
+colorFrom: blue
+colorTo: indigo
+sdk: docker
+app_port: 7860
+pinned: false
+license: mit
+---
+
+# SQL Error Classifier — Inference
+
+Classify student SQL mistakes using a **CodeBERT cross-encoder** model.
+
+Model: [nishu08/sql-codebert-classifier](https://huggingface.co/nishu08/sql-codebert-classifier)
+
+## Inputs
+
+- **Question** — what the student was asked
+- **Schema** — available tables/columns
+- **Student SQL** — student's submission
+- **Correct SQL** — reference answer
+
+## Labels
+
+`JOIN_ERROR`, `AGGREGATION_ERROR`, `FILTER_ERROR`, `WINDOW_FUNCTION_ERROR`,
+`SUBQUERY_ERROR`, `NULL_HANDLING_ERROR`, `PERFORMANCE_ERROR`, `LOGICAL_ERROR`, `SYNTAX_ERROR`
+
+## Secrets (optional)
+
+- `SPACE_MODEL_ID` — override model repo (default: `nishu08/sql-codebert-classifier`)

README_TRAIN_SPACE.md

@@ -0,0 +1,45 @@
+---
+title: SQL Error Classifier Training
+emoji: 🧠
+colorFrom: blue
+colorTo: green
+sdk: docker
+app_port: 7860
+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,98 @@
+"""CodeBERT inference Gradio app for Hugging Face Spaces."""
+
+from __future__ import annotations
+
+import os
+
+import gradio as gr
+
+from src.hf_predict_codebert import CodeBERTSQLErrorClassifier
+
+MODEL_ID = os.getenv("SPACE_MODEL_ID", "nishu08/sql-codebert-classifier")
+
+try:
+    clf = CodeBERTSQLErrorClassifier(MODEL_ID)
+    model_status = f"Model loaded: **{MODEL_ID}**"
+except Exception as exc:
+    clf = None
+    model_status = (
+        f"Could not load `{MODEL_ID}`: {exc}\n\n"
+        "Push your model first: `python scripts/push_codebert_to_hub.py --repo-id nishu08/sql-codebert-classifier`"
+    )
+
+EXAMPLES = [
+    [
+        "What is the average score of students in each department?",
+        "students(id, name, score, department_id) | departments(id, name)",
+        "SELECT department_id, SUM(score) FROM students GROUP BY department_id",
+        "SELECT department_id, AVG(score) FROM students GROUP BY department_id",
+        0.5,
+    ],
+    [
+        "Find students who have not provided an email address.",
+        "students(id, name, email, phone)",
+        "SELECT name FROM students WHERE email = NULL",
+        "SELECT name FROM students WHERE email IS NULL",
+        0.5,
+    ],
+    [
+        "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 JOIN departments",
+        "SELECT students.name, departments.name FROM students INNER JOIN departments ON students.department_id = departments.id",
+        0.5,
+    ],
+]
+
+
+def classify(question, schema, student_sql, correct_sql, threshold):
+    if clf is None:
+        return "Model not loaded.", ""
+    result = clf.predict(
+        question=question.strip(),
+        schema=schema.strip(),
+        student_sql=student_sql.strip(),
+        correct_sql=correct_sql.strip(),
+        threshold=threshold,
+    )
+    summary = (
+        f"### {result['primary_label']}\n"
+        f"Confidence: **{result['primary_confidence']:.1%}**\n\n"
+        f"**Active labels:** {', '.join(result['error_labels']) or 'none'}"
+    )
+    probs = "\n".join(
+        f"- **{k}**: {v:.1%}" for k, v in result["probabilities"].items()
+    )
+    return summary, probs
+
+
+with gr.Blocks(title="SQL Error Classifier", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        f"""
+        # SQL Error Classifier
+        CodeBERT cross-encoder — classifies student SQL mistakes into error categories.
+
+        {model_status}
+        """
+    )
+    with gr.Row():
+        with gr.Column():
+            question = gr.Textbox(label="Question", lines=2)
+            schema = gr.Textbox(label="Schema", lines=2)
+            student_sql = gr.Textbox(label="Student SQL", lines=4)
+            correct_sql = gr.Textbox(label="Correct SQL", lines=4)
+            threshold = gr.Slider(0.1, 0.9, value=0.5, step=0.05, label="Threshold")
+            btn = gr.Button("Classify", variant="primary")
+        with gr.Column():
+            prediction = gr.Markdown(label="Prediction")
+            probabilities = gr.Markdown(label="Probabilities")
+
+    gr.Examples(examples=EXAMPLES, inputs=[question, schema, student_sql, correct_sql, threshold])
+    btn.click(
+        classify,
+        [question, schema, student_sql, correct_sql, threshold],
+        [prediction, probabilities],
+    )
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)

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"])
+```

requirements.txt

@@ -0,0 +1,12 @@
+numpy>=1.24.0
+pandas>=2.0.0
+scikit-learn>=1.3.0
+pyarrow>=14.0.0
+tqdm>=4.66.0
+pyyaml>=6.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,<6.0.0

scripts/deploy_inference_space.sh

@@ -0,0 +1,76 @@
+#!/usr/bin/env bash
+# Deploy CodeBERT inference Space to Hugging Face.
+# Usage: ./scripts/deploy_inference_space.sh [YOUR_USERNAME/space-name]
+set -euo pipefail
+
+SPACE_SUFFIX="${1:-sql-error-classifier}"
+MODEL_REPO="${2:-nishu08/sql-codebert-classifier}"
+TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}"
+
+if [[ -z "${TOKEN}" ]]; then
+  echo "Set HF_TOKEN before deploying."
+  exit 1
+fi
+
+if [[ "${SPACE_SUFFIX}" == */* ]]; then
+  SPACE_ID="${SPACE_SUFFIX}"
+else
+  USERNAME=$(python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+print(api.whoami()["name"])
+PY
+)
+  SPACE_ID="${USERNAME}/${SPACE_SUFFIX}"
+fi
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+WORKDIR=$(mktemp -d)
+trap 'rm -rf "${WORKDIR}"' EXIT
+
+echo "==> Preparing inference Space files..."
+rsync -a \
+  --exclude '.venv' \
+  --exclude 'models' \
+  --exclude 'data' \
+  --exclude '__pycache__' \
+  --exclude '.git' \
+  "${ROOT}/" "${WORKDIR}/"
+
+cp "${ROOT}/README_INFERENCE_SPACE.md" "${WORKDIR}/README.md"
+cp "${ROOT}/Dockerfile.inference" "${WORKDIR}/Dockerfile"
+
+# Bake default model id into Dockerfile
+sed -i.bak "s|ENV SPACE_MODEL_ID=.*|ENV SPACE_MODEL_ID=${MODEL_REPO}|" "${WORKDIR}/Dockerfile"
+rm -f "${WORKDIR}/Dockerfile.bak"
+
+echo "==> Verifying model repo exists: ${MODEL_REPO}"
+python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+api.model_info("${MODEL_REPO}")
+print("Model found.")
+PY
+
+echo "==> Creating Space: ${SPACE_ID}"
+python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+api.create_repo("${SPACE_ID}", repo_type="space", space_sdk="docker", exist_ok=True)
+PY
+
+echo "==> Uploading..."
+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 inference Space",
+)
+PY
+
+echo ""
+echo "Done: https://huggingface.co/spaces/${SPACE_ID}"
+echo "CPU hardware is sufficient for inference."

scripts/deploy_train_space.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+# Deploy CodeBERT training Space to Hugging Face.
+# Usage: ./scripts/deploy_train_space.sh [YOUR_USERNAME/space-name]
+set -euo pipefail
+
+SPACE_SUFFIX="${1:-sql-error-classifier-train}"
+TOKEN="${HF_TOKEN:-${HUGGING_FACE_HUB_TOKEN:-}}"
+
+if [[ -z "${TOKEN}" ]]; then
+  echo "Set HF_TOKEN before deploying."
+  exit 1
+fi
+
+# Resolve full space id — must use an account you own
+if [[ "${SPACE_SUFFIX}" == */* ]]; then
+  SPACE_ID="${SPACE_SUFFIX}"
+else
+  USERNAME=$(python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+print(api.whoami()["name"])
+PY
+)
+  SPACE_ID="${USERNAME}/${SPACE_SUFFIX}"
+fi
+
+ROOT="$(cd "$(dirname "$0")/.." && pwd)"
+WORKDIR=$(mktemp -d)
+trap 'rm -rf "${WORKDIR}"' EXIT
+
+echo "==> Preparing Space files..."
+rsync -a \
+  --exclude '.venv' \
+  --exclude 'models' \
+  --exclude '__pycache__' \
+  --exclude '.git' \
+  --exclude '*.joblib' \
+  "${ROOT}/" "${WORKDIR}/"
+
+cp "${ROOT}/README_TRAIN_SPACE.md" "${WORKDIR}/README.md"
+
+echo "==> Creating Space: ${SPACE_ID}"
+python - <<PY
+from huggingface_hub import HfApi
+api = HfApi(token="${TOKEN}")
+api.create_repo("${SPACE_ID}", repo_type="space", space_sdk="docker", exist_ok=True)
+PY
+
+echo "==> Uploading..."
+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 "Set Hardware → GPU t4-small, Secrets → HF_TOKEN"

scripts/push_codebert_to_hub.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python3
+"""Push a locally trained CodeBERT model to Hugging Face Hub."""
+
+from __future__ import annotations
+
+import argparse
+import json
+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.codebert_labels import load_codebert_labels
+
+DEFAULT_MODEL_DIR = PROJECT_ROOT / "models" / "codebert-cross-encoder"
+MODEL_CARD = PROJECT_ROOT / "hub" / "CODEBERT_MODEL_CARD.md"
+
+
+def push(
+    model_dir: Path = DEFAULT_MODEL_DIR,
+    repo_id: str = "",
+    private: bool = False,
+    token: str | None = None,
+) -> str:
+    if not model_dir.exists():
+        raise FileNotFoundError(f"Model not found at {model_dir}. Train first.")
+
+    if not repo_id:
+        raise ValueError("--repo-id required, e.g. nishu08/sql-codebert-classifier")
+
+    token = token or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
+    if not token:
+        raise ValueError("Set HF_TOKEN environment variable")
+
+    # Ensure label config exists for inference
+    label_config = model_dir / "label_config.json"
+    if not label_config.exists():
+        with open(label_config, "w") as f:
+            json.dump(
+                {
+                    "labels": load_codebert_labels(),
+                    "model_name": "microsoft/codebert-base",
+                    "architecture": "codebert-cross-encoder",
+                    "threshold": 0.5,
+                    "max_length": 512,
+                },
+                f,
+                indent=2,
+            )
+
+    api = HfApi(token=token)
+    create_repo(repo_id, repo_type="model", private=private, exist_ok=True, token=token)
+
+    print(f"Uploading {model_dir} → {repo_id} ...")
+    api.upload_folder(
+        folder_path=str(model_dir),
+        repo_id=repo_id,
+        repo_type="model",
+        token=token,
+        commit_message="Upload CodeBERT SQL error classifier",
+    )
+
+    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 CodeBERT model to HF Hub")
+    parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR)
+    parser.add_argument("--repo-id", type=str, required=True)
+    parser.add_argument("--private", action="store_true")
+    parser.add_argument("--token", type=str, default=None)
+    args = parser.parse_args()
+    push(
+        model_dir=args.model_dir,
+        repo_id=args.repo_id,
+        private=args.private,
+        token=args.token,
+    )
+
+
+if __name__ == "__main__":
+    main()

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/device_utils.py

@@ -0,0 +1,13 @@
+"""Pick the best available torch device."""
+
+from __future__ import annotations
+
+import torch
+
+
+def get_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"

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,177 @@
+"""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.device_utils import get_device
+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"
+
+
+def _is_hub_id(model_dir: Union[str, Path]) -> bool:
+    text = str(model_dir)
+    local = Path(text)
+    return "/" in text and not local.exists()
+
+
+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.hub_id = str(model_dir) if _is_hub_id(model_dir) else None
+        self.model_dir = Path(model_dir) if not self.hub_id else None
+        self.threshold = threshold
+        # MPS inference can be flaky for some ops; CPU is reliable on Mac
+        self.device = device or (
+            "cuda" if torch.cuda.is_available() else "cpu"
+        )
+        model_ref = self.hub_id or str(self.model_dir)
+
+        if self.hub_id:
+            self.label_list = load_codebert_labels()
+            self.max_length = 512
+        else:
+            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(model_ref)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_ref
+        ).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,231 @@
+"""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.device_utils import get_device
+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):,}")
+    device = get_device()
+    use_fp16 = fp16 and device == "cuda"
+    print(f"Device: {device} | fp16: {use_fp16}")
+
+    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=use_fp16,
+        use_mps_device=(device == "mps"),
+        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/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,
+}

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(server_name="0.0.0.0", server_port=7860)