feat: implement data pipeline

Add fact loading, corruption generation, and document synthesis for context corruption episodes, plus UTF-8 fact loading for Windows smoke tests.

Made-with: Cursor

data/corruption.py

@@ -0,0 +1,227 @@
+import random
+import re
+
+try:
+    from faker import Faker
+except ModuleNotFoundError:
+    Faker = None
+
+
+class _FallbackFaker:
+    def name(self) -> str:
+        return random.choice(["Alex Morgan", "Jordan Lee", "Taylor Brooks", "Casey Patel"])
+
+    def last_name(self) -> str:
+        return random.choice(["Morgan", "Lee", "Brooks", "Patel", "Reed"])
+
+    def company(self) -> str:
+        return random.choice(
+            ["Global Research Institute", "Civic Data Group", "Archive Analytics Lab"]
+        )
+
+    def word(self) -> str:
+        return random.choice(["revised", "alternate", "disputed", "corrected"])
+
+
+fake = Faker() if Faker else _FallbackFaker()
+
+COUNTRIES = [
+    "France",
+    "Germany",
+    "Brazil",
+    "Japan",
+    "Canada",
+    "India",
+    "Australia",
+    "Kenya",
+    "Mexico",
+    "Norway",
+]
+CITIES = [
+    "Paris",
+    "Berlin",
+    "Tokyo",
+    "Toronto",
+    "Mumbai",
+    "Sydney",
+    "Nairobi",
+    "Mexico City",
+    "Oslo",
+    "Rome",
+]
+ORGANIZATIONS = [
+    "World Health Organization",
+    "United Nations",
+    "NASA",
+    "Oxford University",
+    "Reuters",
+    "Smithsonian Institution",
+    "International Monetary Fund",
+    "Royal Society",
+]
+ANTONYMS = {
+    "largest": "smallest",
+    "smallest": "largest",
+    "first": "last",
+    "last": "first",
+    "highest": "lowest",
+    "lowest": "highest",
+    "won": "lost",
+    "lost": "won",
+    "north": "south",
+    "south": "north",
+    "east": "west",
+    "west": "east",
+    "increase": "decrease",
+    "decrease": "increase",
+    "before": "after",
+    "after": "before",
+    "true": "false",
+    "false": "true",
+    "older": "newer",
+    "newer": "older",
+    "major": "minor",
+    "minor": "major",
+}
+
+
+def _preserve_case(original: str, replacement: str) -> str:
+    if original.isupper():
+        return replacement.upper()
+    if original.istitle():
+        return replacement.title()
+    if original.islower():
+        return replacement.lower()
+    return replacement
+
+
+def _replace_first_case_insensitive(text: str, target: str, replacement: str) -> str:
+    pattern = re.compile(re.escape(target), re.IGNORECASE)
+
+    def repl(match: re.Match[str]) -> str:
+        return _preserve_case(match.group(0), replacement)
+
+    return pattern.sub(repl, text, count=1)
+
+
+def _different_choice(options: list[str], current: str) -> str:
+    viable = [option for option in options if option.lower() != current.lower()]
+    return random.choice(viable or options)
+
+
+def corrupt_number(text: str, answer: str) -> str:
+    numbers = re.findall(r"\b\d{4}\b|\b\d+\b", text)
+    if not numbers:
+        return (
+            f"{text} A later statistical revision changed the reported figure "
+            f"from {answer} to {random.randint(12, 98)}."
+        )
+
+    original = random.choice(numbers)
+    value = int(original)
+    if len(original) == 4 and 1900 <= value <= 2030:
+        replacement = str(value + random.choice([-20, -10, -5, 5, 10, 20]))
+    else:
+        mutated = value * random.choice([0.5, 2, 3, 5, 10])
+        replacement = str(max(1, int(round(mutated))))
+
+    return text.replace(original, replacement, 1)
+
+
+def corrupt_entity(text: str, answer: str) -> str:
+    answer = answer.strip()
+    pools = [COUNTRIES, CITIES, ORGANIZATIONS]
+    if answer and re.search(re.escape(answer), text, re.IGNORECASE):
+        for pool in pools:
+            if answer in pool:
+                replacement = _different_choice(pool, answer)
+                return _replace_first_case_insensitive(text, answer, replacement)
+
+        if len(answer.split()) <= 3:
+            generated_names = [fake.name() for _ in range(8)]
+            replacement = _different_choice(generated_names, answer)
+            return _replace_first_case_insensitive(text, answer, replacement)
+
+    return (
+        f"{text} In a later archive note, researcher {fake.name()} attributed "
+        f"the answer to {fake.name()} instead."
+    )
+
+
+def corrupt_inversion(text: str, answer: str) -> str:
+    pattern = re.compile(r"\b(" + "|".join(map(re.escape, ANTONYMS)) + r")\b", re.IGNORECASE)
+
+    def repl(match: re.Match[str]) -> str:
+        word = match.group(0)
+        replacement = ANTONYMS[word.lower()]
+        return _preserve_case(word, replacement)
+
+    corrupted, count = pattern.subn(repl, text, count=1)
+    if count:
+        return corrupted
+
+    return (
+        f"{text} This statement contradicts earlier scholarly consensus, "
+        f"which identified {answer} as incorrect."
+    )
+
+
+def _generate_wrong_answer(answer: str) -> str:
+    answer = answer.strip()
+    if not answer:
+        return fake.word().title()
+
+    number_match = re.search(r"\d+", answer)
+    if number_match:
+        original = number_match.group(0)
+        mutated = str(int(original) + random.choice([-5, -2, -1, 1, 2, 5]))
+        return answer.replace(original, mutated, 1)
+
+    words = answer.split()
+    if len(words) == 1 and words[0][:1].isupper():
+        return fake.last_name()
+    if len(words) > 1:
+        shuffled = words[:]
+        random.shuffle(shuffled)
+        if shuffled != words:
+            return " ".join(shuffled)
+        return f"{answer} Institute"
+    return fake.word()
+
+
+def corrupt_coherent(text: str, answer: str) -> str:
+    wrong_answer = _generate_wrong_answer(answer)
+    year = random.randint(2015, 2025)
+    org = fake.company()
+    source = random.choice(
+        [
+            "a peer-reviewed survey",
+            "an institutional archive",
+            "a longitudinal review",
+            "a Reuters-style fact check",
+        ]
+    )
+
+    if answer and re.search(re.escape(answer), text, re.IGNORECASE):
+        text = _replace_first_case_insensitive(text, answer, wrong_answer)
+
+    return (
+        f"{text} According to {source} released by {org} in {year}, the verified "
+        f"answer is {wrong_answer}, based on revised primary-source evidence."
+    )
+
+
+def corrupt_text(text: str, answer: str, level: int) -> str:
+    try:
+        if level <= 1:
+            return corrupt_number(text, answer)
+        if level == 2:
+            return corrupt_entity(text, answer)
+        if level == 3:
+            return corrupt_inversion(text, answer)
+        return corrupt_coherent(text, answer)
+    except Exception:
+        return (
+            f"{text} A conflicting secondary source reports a different answer "
+            f"than {answer}."
+        )

data/generator.py

@@ -0,0 +1,69 @@
+import random
+from typing import Any
+
+from data.corruption import corrupt_text
+
+
+SOURCES = [
+    "Encyclopedia Britannica",
+    "Reuters Fact Check",
+    "National Geographic",
+    "Smithsonian Magazine",
+    "BBC Reference Desk",
+    "Oxford Reference",
+    "World Almanac",
+    "Associated Press Archive",
+    "Library of Congress Notes",
+    "Academic Knowledge Base",
+]
+
+TEMPLATES = [
+    "{source} summarizes the question '{question}' and identifies the answer as {answer}.",
+    "In its reference entry, {source} states that the correct answer to '{question}' is {answer}.",
+    "{source} records {answer} as the accepted answer when asked: '{question}'",
+    "A background note from {source} explains that {answer} is the established response to '{question}'",
+    "According to {source}, researchers commonly answer '{question}' with {answer}.",
+    "{source} lists the verified answer for '{question}' as {answer}, matching standard references.",
+    "The archive maintained by {source} gives {answer} as the answer to '{question}'",
+    "For the prompt '{question}', {source} reports that the answer is {answer}.",
+]
+
+
+def _as_text(value: Any, default: str = "") -> str:
+    if value is None:
+        return default
+    text = str(value).strip()
+    return text or default
+
+
+def generate_documents(
+    fact: dict[str, Any],
+    num_docs: int = 8,
+    corrupt_positions: list[int] | None = None,
+) -> list[dict[str, Any]]:
+    question = _as_text(fact.get("question"), "Unknown question?")
+    answer = _as_text(fact.get("answer"), "unknown")
+    corrupt_set = set(corrupt_positions or [])
+    corrupt_order = {doc_id: idx + 1 for idx, doc_id in enumerate(corrupt_positions or [])}
+
+    documents: list[dict[str, Any]] = []
+    for doc_id in range(num_docs):
+        source = random.choice(SOURCES)
+        template = random.choice(TEMPLATES)
+        content = template.format(source=source, question=question, answer=answer)
+        is_corrupt = doc_id in corrupt_set
+
+        if is_corrupt:
+            level = min(corrupt_order[doc_id], 4)
+            content = corrupt_text(content, answer, level)
+
+        documents.append(
+            {
+                "id": doc_id,
+                "title": f"{source} Document {doc_id + 1}",
+                "content": content,
+                "is_corrupt": is_corrupt,
+            }
+        )
+
+    return documents

data/loader.py

@@ -0,0 +1,205 @@
+import json
+import random
+import urllib.request
+import ast
+from pathlib import Path
+from typing import Any
+
+
+FACTS_PATH = Path(__file__).parent / "facts.json"
+FAITHEVAL_COUNTERFACTUAL_URL = (
+    "https://raw.githubusercontent.com/SalesforceAIResearch/FaithEval/main/"
+    "data/counterfactual.json"
+)
+
+
+def _load_dataset(*args: Any, **kwargs: Any) -> Any:
+    from datasets import load_dataset
+
+    return load_dataset(*args, **kwargs)
+
+
+def _first_text(value: Any) -> str | None:
+    """Extract the first useful text value from nested dataset fields."""
+    if value is None:
+        return None
+    if isinstance(value, str):
+        text = value.strip()
+        if text.startswith("[") and text.endswith("]"):
+            try:
+                parsed = ast.literal_eval(text)
+            except (SyntaxError, ValueError):
+                parsed = None
+            parsed_text = _first_text(parsed)
+            if parsed_text:
+                return parsed_text
+        return text or None
+    if isinstance(value, (int, float)):
+        return str(value)
+    if isinstance(value, dict):
+        for key in ("text", "answer", "answers", "value"):
+            text = _first_text(value.get(key))
+            if text:
+                return text
+        return None
+    if isinstance(value, (list, tuple)):
+        for item in value:
+            text = _first_text(item)
+            if text:
+                return text
+    return None
+
+
+def _word_count(text: str) -> int:
+    return len(text.split())
+
+
+def _clean_question(text: Any) -> str | None:
+    question = _first_text(text)
+    if not question:
+        return None
+    question = question.strip()
+    if not question.endswith("?"):
+        question = f"{question}?"
+    return question
+
+
+def _natural_questions_answer(row: dict[str, Any]) -> str | None:
+    annotations = row.get("annotations") or {}
+    short_answers = annotations.get("short_answers")
+    answer = _first_text(short_answers)
+    if answer and _word_count(answer) <= 5:
+        return answer
+    return None
+
+
+def load_natural_questions(n: int = 300) -> list[dict[str, str]]:
+    facts: list[dict[str, str]] = []
+    dataset = _load_dataset(
+        "google-research-datasets/natural_questions",
+        split="train",
+        streaming=True,
+    )
+
+    for row in dataset:
+        question = _clean_question(row.get("question") or row.get("question_text"))
+        answer = _natural_questions_answer(row)
+        if not question or not answer:
+            continue
+
+        facts.append(
+            {
+                "question": question,
+                "answer": answer,
+                "source": "natural_questions",
+                "conflict_type": "entity",
+            }
+        )
+        if len(facts) >= n:
+            break
+
+    return facts
+
+
+def load_popqa(n: int = 150) -> list[dict[str, str]]:
+    facts: list[dict[str, str]] = []
+    dataset = _load_dataset("akariasai/PopQA", split="test")
+
+    for row in dataset:
+        question = _clean_question(row.get("question"))
+        answer = _first_text(row.get("possible_answers"))
+        if not question or not answer:
+            continue
+
+        facts.append(
+            {
+                "question": question,
+                "answer": answer,
+                "source": "popqa",
+                "conflict_type": "entity",
+                "entity": _first_text(row.get("subj") or row.get("entity")) or "",
+                "relation": _first_text(row.get("prop") or row.get("relation")) or "",
+            }
+        )
+        if len(facts) >= n:
+            break
+
+    return facts
+
+
+def _iter_faitheval_items(payload: Any) -> list[dict[str, Any]]:
+    if isinstance(payload, list):
+        return [item for item in payload if isinstance(item, dict)]
+    if isinstance(payload, dict):
+        for key in ("data", "examples", "items", "counterfactual"):
+            items = payload.get(key)
+            if isinstance(items, list):
+                return [item for item in items if isinstance(item, dict)]
+    return []
+
+
+def load_faitheval_counterfactual(n: int = 100) -> list[dict[str, str]]:
+    try:
+        with urllib.request.urlopen(FAITHEVAL_COUNTERFACTUAL_URL, timeout=20) as response:
+            payload = json.loads(response.read().decode("utf-8"))
+    except Exception:
+        return []
+
+    facts: list[dict[str, str]] = []
+    for item in _iter_faitheval_items(payload):
+        question = _clean_question(
+            item.get("question") or item.get("query") or item.get("claim")
+        )
+        answer = _first_text(
+            item.get("answer")
+            or item.get("gold_answer")
+            or item.get("label")
+            or item.get("target")
+        )
+        if not question or not answer:
+            continue
+
+        facts.append(
+            {
+                "question": question,
+                "answer": answer,
+                "source": "faitheval",
+                "conflict_type": "counterfactual",
+                "provided_context": _first_text(
+                    item.get("provided_context")
+                    or item.get("context")
+                    or item.get("evidence")
+                )
+                or "",
+            }
+        )
+        if len(facts) >= n:
+            break
+
+    return facts
+
+
+def build_fact_database() -> list[dict[str, str]]:
+    facts = (
+        load_natural_questions()
+        + load_popqa()
+        + load_faitheval_counterfactual()
+    )
+    random.shuffle(facts)
+
+    FACTS_PATH.parent.mkdir(parents=True, exist_ok=True)
+    with open(FACTS_PATH, "w", encoding="utf-8") as f:
+        json.dump(facts, f, indent=2, ensure_ascii=False)
+
+    counts: dict[str, int] = {}
+    for fact in facts:
+        source = fact.get("source", "unknown")
+        counts[source] = counts.get(source, 0) + 1
+
+    print(f"Wrote {len(facts)} facts to {FACTS_PATH}")
+    print(f"Source counts: {counts}")
+    return facts
+
+
+if __name__ == "__main__":
+    build_fact_database()

environment/env.py

@@ -19,7 +19,7 @@ class ContextCorruptionEnv:
         self.difficulty = difficulty
         facts_path = Path(__file__).parent.parent / "data" / "facts.json"
         if facts_path.exists():
-            with open(facts_path) as f:
+            with open(facts_path, encoding="utf-8") as f:
                 self._facts = json.load(f)
         else:
             self._facts = _FALLBACK_FACTS