Siddh12334's picture
chore: add CLAUDE.md, README, update gitignore
c12488d
|
Raw
History Blame
17.3 kB

ContextCorruption-Env β€” Engineering Spec

OpenEnv Hackathon | Meta Γ— HuggingFace Γ— PyTorch Team: Siddh + Teammate | Deadline: Tomorrow 5pm | Budget: $60 HF Credits


Division of Ownership

Owner Scope
Siddh environment/ (env, reward, actions, server) + training/
Teammate data/ (loader, corruption, generator, facts.json)
Both Integration smoke test (Tomorrow 9am), Docker + HF deploy (Tomorrow 11am)

Repo Layout

context-corruption-env/
β”œβ”€β”€ CLAUDE.md
β”œβ”€β”€ README.md
β”œβ”€β”€ openenv.yaml
β”œβ”€β”€ Dockerfile
β”œβ”€β”€ requirements.txt
β”œβ”€β”€ environment/          # Siddh
β”‚   β”œβ”€β”€ __init__.py
β”‚   β”œβ”€β”€ actions.py        # Pydantic schemas
β”‚   β”œβ”€β”€ reward.py         # scoring logic
β”‚   β”œβ”€β”€ env.py            # OpenEnv Environment subclass
β”‚   └── server.py         # FastAPI app via OpenEnv helper
β”œβ”€β”€ data/                 # Teammate
β”‚   β”œβ”€β”€ loader.py
β”‚   β”œβ”€β”€ corruption.py
β”‚   β”œβ”€β”€ generator.py
β”‚   └── facts.json        # generated artifact, 500+ QA pairs
β”œβ”€β”€ training/
β”‚   β”œβ”€β”€ train_grpo.py
β”‚   └── ContextCorruption_GRPO.ipynb
β”œβ”€β”€ eval/
β”‚   └── baseline_eval.py
└── assets/
    β”œβ”€β”€ reward_curve.png
    └── loss_curve.png

Bootstrap (both, do this first)

# Create repo on GitHub: Public, Python .gitignore, MIT license
git clone https://github.com/YOUR_USERNAME/context-corruption-env.git
cd context-corruption-env

mkdir -p environment data training eval assets
touch environment/__init__.py
touch environment/{actions,reward,env,server}.py
touch data/{loader,corruption,generator}.py
touch training/train_grpo.py eval/baseline_eval.py
touch requirements.txt openenv.yaml Dockerfile

git add . && git commit -m "feat: initial structure" && git push origin main

# Siddh's branch
git checkout -b feat/environment && git push origin feat/environment

# Teammate's branch
git checkout -b feat/data-pipeline && git push origin feat/data-pipeline

# Shared venv
python -m venv venv && source venv/bin/activate
pip install openenv fastapi uvicorn websockets pydantic \
            datasets transformers torch trl unsloth \
            wandb faker python-dotenv
pip freeze > requirements.txt

environment/actions.py

Purpose: Define the shared data contracts. Everything else imports from here.

ActionType enum β€” four string-valued variants:

  • read_doc β€” read a document by index (costs budget, no state change)
  • flag_suspicious β€” mark a doc as potentially corrupted
  • unflag_doc β€” remove a flag
  • submit_answer β€” end the episode with a final answer

ContextCorruptionAction (Pydantic BaseModel):

  • action_type: ActionType
  • doc_id: Optional[int] β€” 0-indexed, only used for doc actions
  • answer: Optional[str] β€” only used on submit
  • confidence: Optional[float] β€” 0.0–1.0, only used on submit; validate range

Document (Pydantic BaseModel):

  • id: int, title: str, content: str
  • is_flagged: bool = False β€” this is the agent's flag, not ground truth

EpisodeObservation (Pydantic BaseModel):

  • question: str
  • documents: list[Document]
  • flagged_ids: list[int]
  • budget_remaining: int
  • turn: int
  • episode_done: bool = False
  • reward: Optional[float] β€” only populated after SUBMIT_ANSWER or budget exhaustion
  • message: Optional[str] β€” optional human-readable status

environment/reward.py

Purpose: Score a completed episode. No LLM judge β€” fully deterministic.

Function signature:

compute_reward(submitted_answer, ground_truth_answer, flagged_ids,
               corrupt_ids, confidence, budget_used, max_budget)
    -> tuple[float, dict]

Scoring breakdown (weights sum to ~1.05 max, floor ~-0.5):

Component Logic Weight
Answer correctness normalize both strings (lowercase, strip punct, collapse whitespace), exact match +0.4
Flag recall len(true_positives) / len(corrupt_ids) β€” fraction of corrupt docs caught +0.3
Precision (no false flags) start at +0.2, subtract 0.1 per false positive, floor at 0 +0.2
Confidence calibration if correct: +0.1 Γ— confidence; if wrong: -0.2 Γ— confidence Β±0.1
Efficiency 0.05 Γ— (1 - budget_used / max_budget) β€” small bonus, don't over-optimise +0.05

Return (round(total, 4), breakdown_dict). The breakdown dict must include all component keys plus "total".

Private helper: _normalize(text: str) -> str β€” lowercase, strip non-word chars, collapse whitespace.


environment/env.py

Purpose: The stateful RL environment. Subclass openenv.Environment.

Class constants:

  • MAX_BUDGET = 12
  • NUM_DOCS = 8
  • DIFFICULTY_LEVELS = [1, 2, 3, 4] (number of corrupt docs per episode)

__init__(self, difficulty=None)

  • Store difficulty (None = random per episode)
  • Load data/facts.json from Path(__file__).parent.parent / "data" / "facts.json"
  • If the file doesn't exist, fall back to a hardcoded single-item list so the env still imports cleanly
  • Call _reset_state()

reset() -> EpisodeObservation

  • Call _reset_state()
  • Sample a random fact
  • Pick n_corrupt: use self.difficulty if set, else random.choice(DIFFICULTY_LEVELS)
  • Sample n_corrupt positions from range(NUM_DOCS) without replacement β†’ store as self._corrupt_ids
  • Call data.generator.generate_documents(fact, num_docs=NUM_DOCS, corrupt_positions=self._corrupt_ids)
  • Store question, ground truth answer, documents on self
  • Return _build_observation()

step(action: ContextCorruptionAction) -> EpisodeObservation

  • If self._done, return observation with message "Episode already done."
  • Increment self._turn and self._budget_used
  • Dispatch on action.action_type:
    • READ_DOC β€” no-op (document content is already in the observation; the budget cost is the point)
    • FLAG_SUSPICIOUS β€” append action.doc_id to self._flagged_ids if not already there
    • UNFLAG_DOC β€” remove action.doc_id from self._flagged_ids if present
    • SUBMIT_ANSWER β€” call compute_reward(...), set self._done = True, store reward + breakdown
  • After dispatch, check if budget_used >= MAX_BUDGET and not yet done β†’ force submission with empty answer, confidence 0.0
  • Return _build_observation(reward=reward)

state() -> dict β€” return all internal state including ground truth (for logging/debug only, never surfaced to the agent via the observation)

_reset_state() β€” zero/clear all instance variables

_build_observation(reward=None, message=None) -> EpisodeObservation β€” build the Pydantic observation from current state; set is_flagged on each Document based on self._flagged_ids


environment/server.py

Purpose: Expose the env over HTTP/WebSocket for TRL to connect to.

Use OpenEnv's create_app helper β€” it handles session management and WebSocket routing automatically. You only need to pass:

  • env_factory: a zero-arg callable that returns a fresh ContextCorruptionEnv()
  • action_model: ContextCorruptionAction
  • observation_model: EpisodeObservation
  • max_concurrent_envs: 64

Assign the result to app so uvicorn can find it. Add a __main__ guard that runs uvicorn on 0.0.0.0:8000.

Note: If the create_app helper name differs in the installed version, check meta-pytorch/OpenEnv GitHub for the current API surface.


data/loader.py

Purpose: Pull QA facts from three sources, merge, shuffle, write facts.json.

Three loaders (implement as separate functions, called by build_fact_database):

  1. load_natural_questions(n=300)

    • Dataset: google-research-datasets/natural_questions, train split, streaming=True
    • Filter: only rows where annotations.short_answers[0].text exists and has ≀5 words
    • Shape each fact as {question, answer, source: "natural_questions", conflict_type: "entity"}
  2. load_popqa(n=150)

    • Dataset: akariasai/PopQA, test split
    • Filter: rows where possible_answers is non-empty
    • Shape: {question, answer: possible_answers[0], source: "popqa", conflict_type: "entity", entity, relation}
  3. load_faitheval_counterfactual(n=100)

    • Source: raw JSON from the SalesforceAIResearch/FaithEval GitHub repo (data/counterfactual.json)
    • Fetch with urllib.request; wrap in try/except and return [] on any failure (URL may be down)
    • Shape: {question, answer, source: "faitheval", conflict_type: "counterfactual", provided_context}

build_fact_database() β€” call all three, concatenate, shuffle, write to data/facts.json. Also callable as __main__.


data/corruption.py

Purpose: Four escalating corruption functions. This is the creative heart of the dataset.

Level 1 β€” corrupt_number(text, answer)

  • Find all integers in text with regex \b\d{4}\b|\b\d+\b
  • If the number looks like a year (4 digits, 1900–2030), shift by a random offset (Β±5, Β±10, Β±20)
  • Otherwise multiply by a random factor (0.5Γ—, 2Γ—, 3Γ—, 5Γ—, 10Γ—)
  • Fallback if no numbers found: append a note claiming the figure was revised

Level 2 β€” corrupt_entity(text, answer)

  • Maintain pools of plausible substitutes by category: countries, cities, person names (use Faker), organizations
  • If the answer appears in text, replace it with a different member of the most-fitting pool
  • Fallback: append a sentence attributing the fact to a Faker-generated person name

Level 3 β€” corrupt_inversion(text, answer)

  • Maintain a hardcoded antonym map: largest↔smallest, first↔last, highest↔lowest, won↔lost, north↔south, etc.
  • Case-preserving replacement (preserve UPPER/Title/lower casing of the matched word)
  • Fallback: append a sentence saying this contradicts earlier scholarly consensus

Level 4 β€” corrupt_coherent(text, answer)

  • Generate a plausible wrong answer via _generate_wrong_answer:
    • If answer contains digits β†’ mutate a number (Β±1, Β±2, Β±5)
    • If single capitalised word β†’ use fake.last_name()
    • Multi-word β†’ shuffle words
  • Insert the wrong answer into the text and wrap it in an authoritative-sounding template citing a fake source + fake org + plausible year
  • This level should read as convincing to a careful human reader

Dispatcher: corrupt_text(text, answer, level: int) -> str β€” route to the right function, catch all exceptions and return a safe fallback.


data/generator.py

Purpose: Wrap a raw QA fact into a list of 8 document dicts, some corrupted.

generate_documents(fact, num_docs=8, corrupt_positions=None) -> list[dict]

For each document index:

  • Pick a random source name (e.g. "Encyclopedia Britannica", "Reuters Fact Check", etc. β€” maintain a pool of ~10)
  • Pick a random sentence template that incorporates {source}, {question}, and {answer}
  • If the index is in corrupt_positions, apply corrupt_text at an escalating level (first corrupt doc = level 1, second = level 2, etc., capping at 4)
  • Return a dict with id, title, content, is_corrupt (ground truth β€” never shown to the agent)

Integration Smoke Test (Tomorrow 9am β€” run together)

from environment.env import ContextCorruptionEnv
from environment.actions import ContextCorruptionAction, ActionType

env = ContextCorruptionEnv(difficulty=2)
obs = env.reset()
assert len(obs.documents) == 8
assert obs.budget_remaining == 12

obs = env.step(ContextCorruptionAction(action_type=ActionType.READ_DOC, doc_id=0))
assert obs.budget_remaining == 11

obs = env.step(ContextCorruptionAction(action_type=ActionType.FLAG_SUSPICIOUS, doc_id=0))
assert 0 in obs.flagged_ids

obs = env.step(ContextCorruptionAction(
    action_type=ActionType.SUBMIT_ANSWER, answer="test", confidence=0.8))
assert obs.episode_done
assert -0.5 <= obs.reward <= 1.05
print("Smoke test passed. State:", env.state())

This must run without errors before starting training.


training/train_grpo.py

Purpose: Fine-tune Qwen2-1.5B-Instruct with GRPO against the live env server.

Steps to implement:

  1. Init WandB: wandb.init(project="context-corruption-env", name="qwen-1.5b-grpo-run1")

  2. Load model with Unsloth:

    • unsloth/Qwen2-1.5B-Instruct, max_seq_length=2048, load_in_4bit=True
    • Apply LoRA: r=16, target q/k/v/o_proj, no dropout, use_gradient_checkpointing="unsloth"
  3. Configure GRPO:

    • num_train_epochs=3
    • per_device_train_batch_size=4, gradient_accumulation_steps=4
    • learning_rate=5e-5, max_completion_length=512
    • num_generations=8 (GRPO group size)
    • report_to="wandb", logging_steps=10, save_steps=50
  4. System prompt for the agent β€” describe the task, list the three tools (read_doc, flag_suspicious, submit_answer), and give the strategy: cross-reference claims, flag inconsistencies, trust parametric knowledge when docs conflict.

  5. env_factory = lambda: EnvClient(base_url=ENV_URL).sync() β€” fill ENV_URL after HF Space is live.

  6. Construct GRPOTrainer(model, config, environment_factory, tokenizer) and call .train().

  7. After training: wandb.finish(). Download reward + loss plots from WandB and save as assets/reward_curve.png and assets/loss_curve.png. Commit them.

Budget guidance: ~200 steps target; minimum 50 steps with visible upward trend is acceptable. Reserve ~$20 of credits for reruns.


eval/baseline_eval.py

Purpose: Establish a pre-training baseline to make the improvement curves meaningful.

Run 100 episodes with a random agent (randomly flags 0–4 docs, submits "unknown" with 0.5 confidence). Print avg/min/max reward and write to eval/baseline_results.json. Run this before training.


Dockerfile

FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
RUN python -c "from data.loader import build_fact_database; build_fact_database()" || true
EXPOSE 7860
CMD ["uvicorn", "environment.server:app", "--host", "0.0.0.0", "--port", "7860"]

openenv.yaml

name: context-corruption-env
version: "1.0.0"
description: >
  OpenEnv environment for training epistemic robustness in LLMs.
  Agents identify correct answers and flag corrupted documents
  in a multi-doc QA setting with verifiable, objective rewards.
author: "Your Team Name"
license: MIT

environment:
  entrypoint: "environment.server:app"
  action_schema: "environment.actions.ContextCorruptionAction"
  observation_schema: "environment.actions.EpisodeObservation"
  max_concurrent_sessions: 64

reward:
  type: "objective"
  range: [-0.5, 1.05]
  components:
    - {name: answer_correctness,      weight: 0.4}
    - {name: corruption_detection,    weight: 0.3}
    - {name: false_positive_penalty,  weight: 0.2}
    - {name: confidence_calibration,  weight: 0.1}

datasets:
  - {name: Natural Questions, url: "https://huggingface.co/datasets/google-research-datasets/natural_questions"}
  - {name: PopQA,             url: "https://huggingface.co/datasets/akariasai/PopQA"}
  - {name: FaithEval,         url: "https://github.com/SalesforceAIResearch/FaithEval"}

citation: |
  @misc{contextcorruption2026,
    title={ContextCorruption-Env: Training Epistemic Robustness in LLMs},
    year={2026},
    note={OpenEnv Hackathon Submission}
  }

HuggingFace Deployment

pip install huggingface_hub
huggingface-cli login
huggingface-cli repo create context-corruption-env --type space --space_sdk docker
git remote add space https://huggingface.co/spaces/YOUR_HF_USERNAME/context-corruption-env
git push space main

Emergency fallback if Space deploy fails: pip install pyngrok && ngrok http 8000


Sync Schedule

Time Checkpoint
Tonight 12am Siddh: env.py skeleton runs. Teammate: facts.json with 100+ entries.
Tonight 3am Full integration attempt β€” env.reset() works with real data.
Tomorrow 9am Smoke test passes. Both push to main.
Tomorrow 11am HF Space live. Training starts.
Tomorrow 2pm Training done. Plots committed. README filled in.
Tomorrow 4pm Final submission check. All links in README.

Emergency Fallbacks

  • FaithEval URL down: Skip it β€” NQ + PopQA gives 450+ facts, enough.
  • Training too slow: 50 steps with visible upward trend beats no training evidence.
  • OpenEnv API changed: Check meta-pytorch/OpenEnv GitHub for the current create_app / Environment import paths.

README Template (fill in after training)

Key sections judges want to see:

  1. The Problem β€” LLMs defer to wrong retrieved docs even when they know the answer. Standard RLHF makes this worse (cite ClashEval NeurIPS 2024, CANOE May 2025, Knowledgeable-R1 June 2025).
  2. What We Built β€” OpenEnv RL environment, objective reward, no LLM judge.
  3. Results table β€” Random baseline vs trained model: avg reward, answer accuracy, corruption detection rate.
  4. Reward + loss curve images.
  5. Links: HF Space, Colab, blog post, WandB run.

OpenEnv Hackathon, April 2026