Fix dream prompt; move train.py outside package

pyproject.toml

@@ -13,9 +13,11 @@ dependencies = [
  "datasets",
  "torch",
  "peft",
+ "scipy",
  "simple-parsing",
  # 4.0 introduced the breaking change of using return_dict=True by default
  "transformers>=4.0.0",
+ "wandb",
 ]
 version = "0.0.1"
 

w2s/train.py → train.py

@@ -18,11 +18,10 @@
 )
 
 import wandb
-
-from .ds_registry import load_and_process_dataset
-from .knn import gather_hiddens, zeta_filter
-from .loss import log_confidence_loss
-from .roc_auc import roc_auc
+from w2s.ds_registry import load_and_process_dataset
+from w2s.knn import gather_hiddens, zeta_filter
+from w2s.loss import log_confidence_loss
+from w2s.roc_auc import roc_auc
 
 
 @dataclass

w2s/ds_registry.py

@@ -253,7 +253,7 @@ def format_dream(ex, rng):
 
  ans = rng.choice(distractors)
 
- txt = f"Q: {ex['question']} A: {ans}"
+ txt = f"{ex['dialogue']}\n\nQ: {ex['question']} A: {ans}"
  return dict(txt=txt, hard_label=hard_label)
 
 
@@ -410,7 +410,10 @@ def format_openbookqa(ex, rng):
  del distractors[letters.index(ex["answerKey"])]
  ans = rng.choice(distractors)
 
- txt = f"Q: {ex['question_stem']}\n\nA: {ans}"
+ choices = [
+ f"{a}) {t}" for a, t in zip(ex["choices"]["label"], ex["choices"]["text"])
+ ]
+ txt = f"Q: {ex['question_stem']}\n\nChoices:\n{'\n'.join(choices)}\n\nAnswer: {ans}"
  return dict(txt=txt, hard_label=hard_label)