Upload jev_toy/data.py with huggingface_hub

jev_toy/data.py

@@ -0,0 +1,118 @@
+"""
+jev_toy/data.py
+
+Turn real public HuggingFace datasets into System One training/eval examples,
+each of the form (state_text, question_text, q_type, target).
+
+We map three proven Jev question types onto real datasets:
+
+  * choice  <- AG News (route a headline/topic into 1 of 4 categories)
+  * noul    <- BoolQ   (yes / no on a reading-comprehension question)
+  * score   <- SST-2   (continuous "sentiment polarity" tag)
+
+This is a toy-scale demonstration: we subsample so a laptop CPU can train in
+minutes. The structure generalizes to the full datasets and to GPU.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import Counter
+
+SPLIT_RE = re.compile(r"[A-Za-z0-9']+|[.,!?;:()\"]")
+
+
+def tokenize(text: str, vocab, max_len: int, oov: int, pad: int):
+    ids = []
+    for tok in SPLIT_RE.findall(text.lower())[:max_len]:
+        ids.append(vocab.get(tok, oov))  # certain words map to a small rand-pool
+    ids = ids[:max_len]
+    mask = [1] * len(ids)
+    ids = ids + [pad] * (max_len - len(ids))
+    mask = mask + [0] * (max_len - len(mask))
+    return ids, mask
+
+
+class Vocabulary:
+    def __init__(self, min_freq=1):
+        self.min_freq = min_freq
+        self.stoi = {}
+        self.itos = {}
+        self.oov = 0
+
+    def build(self, texts: list[str]):
+        c = Counter()
+        for t in texts:
+            c.update(SPLIT_RE.findall(t.lower()))
+        words = [w for w, n in c.items() if n >= self.min_freq]
+        words.sort()
+        self.stoi = {w: i + 2 for i, w in enumerate(words)}  # 0=pad, 1=oov
+        self.itos = {i: w for w, i in self.stoi.items()}
+        self.oov = 1
+        return self
+
+    def __len__(self):
+        return len(self.stoi) + 2
+
+
+def make_agnews(df) -> list[dict]:
+    labels = ["World", "Sports", "Business", "Sci/Tech"]
+    out = []
+    for txt, lab in zip(df["text"], df["label"]):
+        out.append({
+            "state": f"News article: {txt[:200]}",
+            "question": "Which topic does this article belong to?",
+            "q_type": "choice",
+            "target": int(lab),
+            "options": labels,
+        })
+    return out
+
+
+def make_boolq(df) -> list[dict]:
+    out = []
+    for passage, q, ans in zip(df["passage"], df["question"], df["answer"]):
+        out.append({
+            "state": f"Passage: {passage[:250]}",
+            "question": f"Q: {q}",
+            "q_type": "noul",
+            "target": int(ans),
+        })
+    return out
+
+
+def make_sst2(df) -> list[dict]:
+    out = []
+    for txt, lab in zip(df["sentence"], df["label"]):
+        out.append({
+            "state": f"Review: {txt[:150]}",
+            "question": "Is this review positive?",
+            "q_type": "noul",      # treat sentiment polarity as a boolean tag
+            "target": int(lab),
+        })
+    return out
+
+
+def build_examples(datasets: dict[str, object], subsample: dict[str, int]) -> tuple[list[dict], Vocabulary]:
+    """datasets: {name: HF-dataset}; subsample: {name: cap}."""
+    all_ex = []
+    # AG News
+    for split in ("train", "test"):
+        df = datasets["agnews"][split]
+        cap = subsample.get("agnews", 2000) if split == "train" else 400
+        all_ex += make_agnews(df.select(list(range(cap))))
+    # BoolQ
+    train = datasets["boolq"]["train"].select(list(range(subsample.get("boolq", 2000))))
+    val = datasets["boolq"]["validation"].select(list(range(400)))
+    all_ex += make_boolq(train) + make_boolq(val)
+    # SST-2
+    sst_train = datasets["sst2"]["train"].select(list(range(subsample.get("sst2", 2000))))
+    sst_val = datasets["sst2"]["validation"].select(list(range(400)))
+    all_ex += make_sst2(sst_train) + make_sst2(sst_val)
+    # vocab from state+question text
+    vocab = Vocabulary().build([e["state"] + " " + e["question"] for e in all_ex])
+    return all_ex, vocab
+
+
+if __name__ == "__main__":
+    print("data module ok")