Fix Transformers batch generation and targeted dataset loading

app.py

@@ -30,10 +30,14 @@ ans = (a + b);"""
 def _load_semantics() -> dict[tuple[str, str], SemanticsSpec]:
     specs: dict[tuple[str, str], SemanticsSpec] = {}
     for formalization in ("S", "K"):
+        standard_split_name = f"{formalization}_Standard_Human_Written"
         standard_split = load_dataset(
-            DATASET_ID,
-            "predstate",
-            split=f"{formalization}_Standard_Human_Written",
+            "parquet",
+            data_files=(
+                f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/"
+                f"predstate/{standard_split_name}-00000-of-00001.parquet"
+            ),
+            split="train",
         )
         standard = standard_split[0]
         specs[(formalization, "Standard")] = SemanticsSpec(
@@ -43,9 +47,12 @@ def _load_semantics() -> dict[tuple[str, str], SemanticsSpec]:
 
         nonstandard_split_name = f"{formalization}_NonStandard_Human_Written"
         nonstandard_split = load_dataset(
-            DATASET_ID,
-            "predstate",
-            split=nonstandard_split_name,
+            "parquet",
+            data_files=(
+                f"https://huggingface.co/datasets/{DATASET_ID}/resolve/main/"
+                f"predstate/{nonstandard_split_name}-00000-of-00001.parquet"
+            ),
+            split="train",
         )
         for intervention in ("KeywordSwap", "KeywordObf"):
             record = next(
@@ -71,7 +78,7 @@ SEMANTICS_SPECS = _load_semantics()
 TOKENIZER = AutoTokenizer.from_pretrained(MODEL_ID)
 MODEL = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
 ).to("cuda")
 MODEL.eval()
 
@@ -99,19 +106,22 @@ def predict(
     model_inputs = TOKENIZER.apply_chat_template(
         messages,
         add_generation_prompt=True,
+        tokenize=True,
         return_tensors="pt",
+        return_dict=True,
     ).to("cuda")
+    input_ids = model_inputs["input_ids"]
 
     with torch.inference_mode():
         generated = MODEL.generate(
-            model_inputs,
+            **model_inputs,
             max_new_tokens=512,
             do_sample=False,
             pad_token_id=TOKENIZER.eos_token_id,
         )
 
     response = TOKENIZER.decode(
-        generated[0, model_inputs.shape[-1] :],
+        generated[0, input_ids.shape[-1] :],
         skip_special_tokens=True,
     ).strip()
     answer = _extract_answer(response)