use hard labels for auroc and acc

underspec/roc_auc.py

@@ -25,6 +25,8 @@ def roc_auc(y_true: Tensor, y_pred: Tensor) -> Tensor:
  )
  if y_true.dim() not in (1, 2):
  raise ValueError("y_true and y_pred should be 1D or 2D tensors")
+ if not ((y_true == 1) | (y_true == 0)).all():
+ raise ValueError("y_true should contain only 0s and 1s")
 
  # Sort y_pred in descending order and get indices
  indices = y_pred.argsort(descending=True, dim=-1)

underspec/sft.py

@@ -85,9 +85,10 @@ def process(examples):
 
  def compute_metrics(eval_pred):
  predictions, labels = map(torch.from_numpy, eval_pred)
+ hard_labels = (labels > 0.5).long()
  return dict(
- accuracy=predictions.argmax(dim=1).eq(labels).float().mean(),
- auroc=roc_auc(labels, predictions[:, 1]),
+ accuracy=predictions.argmax(dim=1).eq(hard_labels).float().mean(),
+ auroc=roc_auc(hard_labels, predictions[:, 1]),
  )
 
  trainer = CustomLossTrainer(