Smoke tests with tiny gpt2, fix CCSReporter

elk/training/ccs_reporter.py

@@ -195,7 +195,11 @@ def loss(
 
  alpha = self.config.supervised_weight
  preds = p0.add(1 - p1).mul(0.5).squeeze(-1)
- bce_loss = bce(preds, labels.type_as(preds))
+ # broadcast the labels, and flatten the predictions
+ # so that both are 1D tensors
+ broadcast_labels = labels.repeat_interleave(preds.shape[1]).float()
+ flattened_preds = preds.cpu().flatten()
+ bce_loss = bce(flattened_preds, broadcast_labels.type_as(flattened_preds))
  loss = alpha * bce_loss + (1 - alpha) * loss
 
  elif self.config.supervised_weight > 0:

tests/test_smoke_elicit.py

@@ -0,0 +1,66 @@
+from pathlib import Path
+
+import pytest
+
+from elk import ExtractionConfig
+from elk.extraction import PromptConfig
+from elk.training import CcsReporterConfig, EigenReporterConfig
+from elk.training.train import train, RunConfig
+
+"""
+TODO: These tests should work with deberta
+but you'll need to make deberta fp32 instead of fp16
+because pytorch cpu doesn't support fp16
+"""
+
+
+def test_smoke_elicit_run_tiny_gpt2_ccs(tmp_path: Path):
+ model_path = "sshleifer/tiny-gpt2"
+ dataset_name = "imdb"
+ config = RunConfig(
+ data=ExtractionConfig(
+ model=model_path,
+ prompts=PromptConfig(dataset=dataset_name, max_examples=[10]),
+ # run on all layers, tiny-gpt only has 2 layers
+ ),
+ net=CcsReporterConfig(),
+ )
+ train(config, tmp_path)
+ # get the files in the tmp_path
+ files: Path = list(tmp_path.iterdir())
+ created_file_names = {file.name for file in files}
+ expected_files = ["cfg.yaml", "metadata.yaml", "lr_models", "reporters", "eval.csv"]
+ for file in expected_files:
+ assert file in created_file_names
+
+
+@pytest.mark.skip(reason="Fix me: EigenReporter crashes with tiny gpt2")
+def test_smoke_elicit_run_tiny_gpt2_eigen(tmp_path: Path):
+ """
+ Currently this test fails with
+ u -= torch.einsum("...ij,...i->...j", V[..., :k, :], proj)
+ V[..., k, :] = F.normalize(u, dim=-1)
+ ~~~~~~~~~ <--- HERE
+
+ u[:] = torch.einsum("...ij,...j->...i", A, V[..., k, :])
+
+ RuntimeError: select(): index 1 out of range for tensor of size [1, 2]
+ at dimension 0
+ """
+ model_path = "sshleifer/tiny-gpt2"
+ dataset_name = "imdb"
+ config = RunConfig(
+ data=ExtractionConfig(
+ model=model_path,
+ prompts=PromptConfig(dataset=dataset_name, max_examples=[10]),
+ # run on all layers, tiny-gpt only has 2 layers
+ ),
+ net=EigenReporterConfig(),
+ )
+ train(config, tmp_path)
+ # get the files in the tmp_path
+ files: Path = list(tmp_path.iterdir())
+ created_file_names = {file.name for file in files}
+ expected_files = ["cfg.yaml", "metadata.yaml", "lr_models", "reporters", "eval.csv"]
+ for file in expected_files:
+ assert file in created_file_names