Code / scripts/ml_training/essay_scoring/training/qwk.py
scripts/ml_training/essay_scoring/training/qwk.py


  1
  """Quadratic weighted kappa utilities."""


  2
  


  3
  from __future__ import annotations


  4
  


  5
  from collections.abc import Callable


  6
  from typing import TYPE_CHECKING


  7
  


  8
  import numpy as np


  9
  from sklearn.metrics import cohen_kappa_score


  10
  


  11
  if TYPE_CHECKING:


  12
      import xgboost as xgb


  13
  


  14
  


  15
  def round_to_half_band(predictions: np.ndarray) -> np.ndarray:


  16
      """Round predictions to the nearest 0.5 band."""


  17
  


  18
      return np.round(predictions * 2.0) / 2.0


  19
  


  20
  


  21
  def clip_bands(predictions: np.ndarray, min_band: float = 1.0, max_band: float = 9.0) -> np.ndarray:


  22
      """Clip predictions to the valid IELTS band range."""


  23
  


  24
      return np.clip(predictions, min_band, max_band)


  25
  


  26
  


  27
  def qwk_score(


  28
      y_true: np.ndarray,


  29
      y_pred: np.ndarray,


  30
      *,


  31
      min_band: float = 1.0,


  32
      max_band: float = 9.0,


  33
  ) -> float:


  34
      """Compute quadratic weighted kappa for predictions."""


  35
  


  36
      rounded = clip_bands(round_to_half_band(y_pred), min_band=min_band, max_band=max_band)


  37
      true_ids = _to_half_band_ids(y_true)


  38
      pred_ids = _to_half_band_ids(rounded)


  39
      return float(cohen_kappa_score(true_ids, pred_ids, weights="quadratic"))


  40
  


  41
  


  42
  def qwk_eval(preds: np.ndarray, dtrain: xgb.DMatrix) -> tuple[str, float]:


  43
      """XGBoost custom evaluation metric for QWK."""


  44
  


  45
      labels = dtrain.get_label()


  46
      return "qwk", qwk_score(labels, preds)


  47
  


  48
  


  49
  def qwk_eval_factory(


  50
      *,


  51
      min_band: float,


  52
      max_band: float,


  53
  ) -> Callable[[np.ndarray, xgb.DMatrix], tuple[str, float]]:


  54
      """Create an XGBoost custom evaluation metric for a specific score scale."""


  55
  


  56
      def qwk_eval_inner(preds: np.ndarray, dtrain: xgb.DMatrix) -> tuple[str, float]:


  57
          labels = dtrain.get_label()


  58
          return "qwk", qwk_score(labels, preds, min_band=min_band, max_band=max_band)


  59
  


  60
      return qwk_eval_inner


  61
  


  62
  


  63
  def _to_half_band_ids(values: np.ndarray) -> np.ndarray:


  64
      """Convert band scores to discrete half-band integer IDs."""


  65
  


  66
      return np.round(values * 2.0).astype(int)