Clean initial commit

.gitignore

@@ -0,0 +1,6 @@
+venv/
+__pycache__/
+*.pyc
+.cache/
+*.egg-info/
+.env

Dockerfile

@@ -0,0 +1,21 @@
+FROM python:3.11-slim
+
+ENV PYTHONUNBUFFERED=1
+ENV HF_DATASETS_CACHE=/app/.cache
+
+WORKDIR /app
+
+# Copy requirements first to leverage Docker layer caching
+COPY requirements.txt .
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the project files
+COPY . .
+
+# Expose HuggingFace Spaces default port
+EXPOSE 7860
+
+# Start the FastAPI server
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,210 @@
+# 🏷️ LabelSense — AI-Assisted Data Labeling QA Environment
+
+> An OpenEnv-compliant environment where an AI agent audits AI-generated labels on medical and legal text, identifies mislabeled examples, flags ambiguous cases, and proposes corrections.
+
+---
+
+## Why This Exists
+
+AI models are increasingly used to auto-label training data at scale. The problem: label quality is inconsistent, and **confidently wrong labels corrupt downstream models silently**. No one catches them until the model is already in production.
+
+LabelSense gives agents a structured environment to practice exactly this — reviewing batches of AI-generated labels, catching errors, and making calibrated decisions about ambiguity. It maps directly to a real MLOps workflow that every team doing data labeling at scale deals with daily.
+
+---
+
+## Environment Overview
+
+The agent is shown examples from real medical and legal datasets, each pre-labeled by a simulated AI labeler (with injected noise). The agent must:
+
+- Decide if each label is **correct**, **wrong**, or **ambiguous**
+- If wrong — propose the **correct label**
+- Signal **confidence** in its decision
+
+The environment scores the agent on accuracy, calibration, and how responsibly it handles uncertainty.
+
+---
+
+## Datasets
+
+| Task | Dataset | Source | Label Type |
+|------|---------|--------|------------|
+| Easy | Medical Question Pairs | `curaihealth/medical_questions_pairs` | Binary: similar / not similar |
+| Medium | Stanford NLI | `snli` | 3-class: entailment / neutral / contradiction |
+| Hard | SCOTUS (LexGLUE) | `coastalcph/lex_glue` (scotus) | 14-class: Supreme Court issue areas |
+
+---
+
+## Action & Observation Space
+
+### Observation
+What the agent sees at each step:
+
+```json
+{
+  "example_id": "task1_042",
+  "task": "easy",
+  "input": {
+    "text1": "Does ibuprofen reduce fever?",
+    "text2": "Can ibuprofen be used to treat high temperature?"
+  },
+  "ai_label": "not_similar",
+  "label_options": ["similar", "not_similar"]
+}
+```
+
+### Action
+What the agent responds with:
+
+```json
+{
+  "example_id": "task1_042",
+  "verdict": "wrong",
+  "proposed_label": "similar",
+  "confidence": 0.91
+}
+```
+
+`verdict` must be one of: `correct`, `wrong`, `ambiguous`
+
+---
+
+## Reward Function
+
+| Agent Action | Condition | Reward |
+|---|---|---|
+| Flags label as wrong | Label is actually wrong | +1.0 |
+| Proposes correct fix | Fix matches gold label | +0.5 bonus |
+| Flags as ambiguous | Example is genuinely ambiguous | +0.7 |
+| Flags as ambiguous | Example is actually clear | -0.3 |
+| Marks wrong label as correct | Misses the error | 0.0 |
+| Proposes wrong fix confidently | High confidence, wrong answer | -0.5 |
+
+Rewards are designed to encourage **calibrated uncertainty** — an agent that admits it doesn't know scores better than one that guesses confidently and gets it wrong.
+
+---
+
+## Tasks
+
+### Task 1 — Easy: Medical Question Pair Similarity
+**Dataset:** `curaihealth/medical_questions_pairs`
+**Label type:** Binary (0 = not similar, 1 = similar)
+**Noise:** ~20% random label flips on clear-cut examples
+**Expected agent score:** 0.75 – 0.90
+**What makes it easy:** Labels are mostly unambiguous. Errors are random, not systematic.
+
+### Task 2 — Medium: Natural Language Inference
+**Dataset:** `snli`
+**Label type:** 3-class (entailment / neutral / contradiction)
+**Noise:** Systematic bias — AI labeler over-predicts "neutral" when uncertain
+**Expected agent score:** 0.55 – 0.75
+**What makes it medium:** Requires understanding sentence-level logic. Neutral vs. contradiction is a common confusion point.
+
+### Task 3 — Hard: SCOTUS Legal Issue Classification
+**Dataset:** `coastalcph/lex_glue` (scotus config)
+**Label type:** 14-class Supreme Court issue areas
+**Noise:** Confident wrong labels on edge cases, near-duplicate category confusion
+**Expected agent score:** 0.30 – 0.55
+**What makes it hard:** 14 overlapping legal categories. AI labeler is confidently wrong, not randomly wrong. Requires legal domain reasoning.
+
+---
+
+## API Endpoints
+
+| Method | Endpoint | Description |
+|--------|----------|-------------|
+| POST | `/reset` | Start a new episode, returns first observation |
+| POST | `/step` | Submit an action, returns next observation + reward |
+| GET | `/state` | Returns current episode state and progress |
+| GET | `/tasks` | Lists all tasks and their action schemas |
+| POST | `/grader` | Returns grader score after episode completes |
+| POST | `/baseline` | Runs baseline inference script, returns scores for all 3 tasks |
+
+---
+
+## Setup & Usage
+
+### Local (Python)
+
+```bash
+# Clone the repo
+git clone https://github.com/yourusername/labelsense-openenv
+cd labelsense-openenv
+
+# Create and activate virtual environment
+python -m venv venv
+source venv/bin/activate  # Windows: venv\Scripts\activate
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Start the API server
+uvicorn main:app --reload
+```
+
+### Docker
+
+```bash
+docker build -t labelsense-env .
+docker run -p 8000:8000 labelsense-env
+```
+
+### Run Baseline
+
+```bash
+export OPENAI_API_KEY=your_key_here
+python baseline.py
+```
+
+---
+
+## OpenEnv Spec Compliance
+
+- Typed Pydantic models for `Observation`, `Action`, `Reward`
+- `reset()` → returns clean initial observation
+- `step(action)` → returns observation, reward, done, info
+- `state()` → returns current episode state
+- `openenv.yaml` with full metadata
+- Validated with `openenv validate`
+
+---
+
+## Baseline Scores
+
+| Task | Model | Cumulative Score (10 steps) |
+|------|-------|-----------------------------|
+| Easy (medical pairs) | llama-3.1-8b-instant | 4.30 |
+| Medium (NLI) | llama-3.1-8b-instant | 2.00 |
+| Hard (SCOTUS) | llama-3.1-8b-instant | 4.90 |
+| **Overall Average** | | **3.73** |
+
+*Scores represent cumulative reward over 10 steps per task. Maximum possible score per task is 10.0 (all correct with fixes). Scores above 0 indicate the agent performs better than random.*
+
+---
+
+## Project Structure
+
+```
+labelsense-openenv/
+├── environment/
+│   ├── env.py          # Core environment — reset(), step(), state()
+│   ├── models.py       # Pydantic models: Observation, Action, Reward
+│   ├── grader.py       # Per-task scoring logic
+│   └── noise.py        # AI labeler noise injection
+├── tasks/
+│   ├── task1_easy.py
+│   ├── task2_medium.py
+│   └── task3_hard.py
+├── data/
+│   └── loader.py       # HuggingFace dataset loading + sampling
+├── main.py             # FastAPI app + all endpoints
+├── baseline.py         # Baseline inference script (OpenAI API)
+├── openenv.yaml        # OpenEnv spec metadata
+├── Dockerfile
+└── README.md
+```
+
+---
+
+## License
+
+MIT

baseline.py

@@ -0,0 +1,150 @@
+import os
+import sys
+import json
+import requests
+from groq import Groq
+
+API_URL = "http://localhost:8000"
+
+def parse_json_response(text: str) -> dict:
+    try:
+        text = text.strip()
+        if text.startswith("```"):
+            lines = text.split("\n")
+            if len(lines) >= 2:
+                # Remove starting and ending markdown fences
+                if lines[-1].strip() == "```":
+                    text = "\n".join(lines[1:-1])
+                else:
+                    text = "\n".join(lines[1:])
+                    
+        parsed = json.loads(text)
+        
+        return {
+            "verdict": parsed.get("verdict", "ambiguous"),
+            "proposed_label": parsed.get("proposed_label"),
+            "confidence": float(parsed.get("confidence", 0.5))
+        }
+    except Exception:
+        # Default back to ambiguous on parse failure
+        return {"verdict": "ambiguous", "proposed_label": None, "confidence": 0.5}
+
+def run_baseline() -> dict:
+    api_key = os.environ.get("GROQ_API_KEY")
+    if not api_key:
+        print("Error: GROQ_API_KEY environment variable is missing.", file=sys.stderr)
+        sys.exit(1)
+        
+    client = Groq(api_key=api_key)
+    model = "llama-3.1-8b-instant"
+    
+    tasks = ["easy", "medium", "hard"]
+    results = {}
+    
+    for task in tasks:
+        session_id = f"baseline_{task}"
+        
+        # 1. Reset Env
+        reset_res = requests.post(f"{API_URL}/reset", json={
+            "task": task,
+            "episode_length": 10,
+            "session_id": session_id
+        })
+        reset_res.raise_for_status()
+        obs = reset_res.json()
+        
+        done = False
+        while not done:
+            # 2. Extract inputs depending on task definition
+            input_fields = obs.get("input", {})
+            if task == "easy":
+                input_text = f"Text 1: {input_fields.get('text1')}\nText 2: {input_fields.get('text2')}"
+            elif task == "medium":
+                input_text = f"Premise: {input_fields.get('premise')}\nHypothesis: {input_fields.get('hypothesis')}"
+            else: # hard
+                raw_text = input_fields.get("text", "")
+                input_text = f"Text: {raw_text[:300]}"  # Truncated to 300 chars
+                
+            ai_label = obs.get("ai_label")
+            label_options = obs.get("label_options")
+            
+            # 3. Create Model Prompt
+            prompt = f"""You are an expert AI auditor verifying labels for a dataset.
+Your task is to review the provided input and decide if the assigned 'AI Label' is correct, wrong, or ambiguous.
+
+Input Examples:
+{input_text}
+
+AI Label: {ai_label}
+Valid Label Options: {label_options}
+
+Instructions:
+Evaluate the AI Label against the Input Examples.
+Respond in pure JSON format only with the following keys:
+- "verdict": purely one of "correct", "wrong", or "ambiguous"
+- "proposed_label": if the verdict is "wrong", provide the correct label from the Valid Label Options as a string. Otherwise, use null.
+- "confidence": a float between 0.0 and 1.0 representing your confidence.
+
+Example of valid response:
+{{"verdict": "wrong", "proposed_label": "1", "confidence": 0.85}}
+"""
+
+            # 4. Invoke LLM
+            chat_completion = client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You output JSON strictly."
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                model=model,
+                temperature=0.0
+            )
+            
+            # 5. Parse output
+            response_text = chat_completion.choices[0].message.content
+            action_dict = parse_json_response(response_text)
+            
+            # 6. Step Env
+            step_payload = {
+                "session_id": session_id,
+                "example_id": obs.get("example_id"),
+                "verdict": action_dict["verdict"],
+                "proposed_label": action_dict.get("proposed_label"),
+                "confidence": action_dict.get("confidence")
+            }
+            
+            step_res = requests.post(f"{API_URL}/step", json=step_payload)
+            step_res.raise_for_status()
+            step_data = step_res.json()
+            
+            done = step_data.get("done", True)
+            if not done:
+                obs = step_data.get("observation", {})
+                
+        # 7. Collect Final Grade
+        grader_res = requests.post(f"{API_URL}/grader", params={"session_id": session_id})
+        grader_res.raise_for_status()
+        final_info = grader_res.json()
+        
+        results[task] = final_info.get("cumulative_score", 0.0)
+        
+    return results
+
+if __name__ == "__main__":
+    print("Running baseline evaluations... (This evaluates the Groq model against the live API)")
+    scores = run_baseline()
+    
+    print("\n--- Baseline Results ---")
+    total_score = 0.0
+    for task, score in scores.items():
+        print(f"Task: {task.capitalize():<10} | Cumulative Score: {score:>5.2f}")
+        total_score += score
+        
+    avg_score = total_score / len(scores) if scores else 0.0
+    print("-" * 35)
+    print(f"Overall Average Score: {avg_score:>5.2f}")

data/loader.py

@@ -0,0 +1,209 @@
+"""
+data/loader.py - HuggingFace dataset loading for OpenEnv Labeling QA.
+
+Loads three classification datasets, samples 150 examples from each,
+and returns them as clean Python lists of dicts with standardized keys.
+
+Dependencies: pip install datasets
+"""
+
+import sys
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+SAMPLE_SIZE = 150
+SEED = 42
+
+# NLI label mapping (int -> str) used by bigbio NLI datasets
+_NLI_LABEL_MAP = {0: "entailment", 1: "neutral", 2: "contradiction"}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _load_hf_dataset(path: str, split: str = "train", name: str = None):
+    """
+    Wrapper around datasets.load_dataset that handles trust_remote_code
+    gracefully across different versions of the `datasets` library.
+    """
+    from datasets import load_dataset
+    import inspect
+
+    kwargs = {"path": path, "split": split}
+    if name is not None:
+        kwargs["name"] = name
+
+    # Only pass trust_remote_code if the installed version supports it
+    sig = inspect.signature(load_dataset)
+    if "trust_remote_code" in sig.parameters:
+        kwargs["trust_remote_code"] = True
+
+    return load_dataset(**kwargs)
+
+
+def _sample(dataset, n: int, seed: int = SEED):
+    """Return a random sample of *n* rows from a HuggingFace Dataset."""
+    if len(dataset) <= n:
+        return dataset
+    return dataset.shuffle(seed=seed).select(range(n))
+
+
+def _safe_str(value) -> str:
+    """Convert a value to a stripped string, handling None gracefully."""
+    if value is None:
+        return ""
+    return str(value).strip()
+
+
+# ---------------------------------------------------------------------------
+# Task 1 - Medical Question Pairs (binary: 0 / 1)
+# Dataset: curaihealth/medical_questions_pairs
+# ---------------------------------------------------------------------------
+
+def load_task1() -> list[dict]:
+    """
+    Load the curaihealth/medical_questions_pairs dataset.
+
+    Returns a list of 150 dicts:
+        {id: str, text1: str, text2: str, gold_label: int}
+    where gold_label is 0 or 1.
+    """
+    try:
+        print("[Task 1] Loading curaihealth/medical_questions_pairs ...")
+        ds = _load_hf_dataset("curaihealth/medical_questions_pairs", split="train")
+
+        sampled = _sample(ds, SAMPLE_SIZE)
+
+        results: list[dict] = []
+        for idx, row in enumerate(sampled):
+            results.append({
+                "id": f"task1_{idx}",
+                "text1": _safe_str(row.get("question_1", row.get("question1", ""))),
+                "text2": _safe_str(row.get("question_2", row.get("question2", ""))),
+                "gold_label": int(row.get("label", 0)),
+            })
+
+        print(f"[Task 1] OK - Loaded {len(results)} examples.")
+        return results
+
+    except Exception as exc:
+        print(f"[Task 1] FAILED - {exc}", file=sys.stderr)
+        raise
+
+
+# ---------------------------------------------------------------------------
+# Task 2 - NLI (3-class: entailment / neutral / contradiction)
+# Dataset: snli
+# ---------------------------------------------------------------------------
+
+def load_task2() -> list[dict]:
+    """
+    Load the Stanford NLI (snli) dataset.
+
+    Filters out unlabeled examples (label == -1), then samples 150.
+
+    Returns a list of 150 dicts:
+        {id: str, premise: str, hypothesis: str, gold_label: str}
+    where gold_label is "entailment", "neutral", or "contradiction".
+    """
+    try:
+        print("[Task 2] Loading snli ...")
+        ds = _load_hf_dataset("snli", split="train")
+
+        # SNLI contains some unlabeled rows marked with label == -1
+        ds = ds.filter(lambda x: x["label"] != -1)
+
+        sampled = _sample(ds, SAMPLE_SIZE)
+
+        results: list[dict] = []
+        for idx, row in enumerate(sampled):
+            label_str = _NLI_LABEL_MAP.get(row["label"], str(row["label"]))
+            results.append({
+                "id": f"task2_{idx}",
+                "premise": _safe_str(row.get("premise", "")),
+                "hypothesis": _safe_str(row.get("hypothesis", "")),
+                "gold_label": label_str,
+            })
+
+        print(f"[Task 2] OK - Loaded {len(results)} examples.")
+        return results
+
+    except Exception as exc:
+        print(f"[Task 2] FAILED - {exc}", file=sys.stderr)
+        raise
+
+
+# ---------------------------------------------------------------------------
+# Task 3 - SCOTUS Legal Classification (14-class: 0-13)
+# Dataset: coastalcph/lex_glue  config="scotus"
+# ---------------------------------------------------------------------------
+
+def load_task3() -> list[dict]:
+    """
+    Load the coastalcph/lex_glue (scotus) dataset.
+
+    Returns a list of 150 dicts:
+        {id: str, text: str, gold_label: int}
+    where gold_label is an int 0-13.
+    """
+    try:
+        print("[Task 3] Loading coastalcph/lex_glue (scotus) ...")
+        ds = _load_hf_dataset("coastalcph/lex_glue", split="train",
+                              name="scotus")
+
+        sampled = _sample(ds, SAMPLE_SIZE)
+
+        results: list[dict] = []
+        for idx, row in enumerate(sampled):
+            results.append({
+                "id": f"task3_{idx}",
+                "text": _safe_str(row.get("text", "")),
+                "gold_label": int(row.get("label", 0)),
+            })
+
+        print(f"[Task 3] OK - Loaded {len(results)} examples.")
+        return results
+
+    except Exception as exc:
+        print(f"[Task 3] FAILED - {exc}", file=sys.stderr)
+        raise
+
+
+# ---------------------------------------------------------------------------
+# Main - quick smoke test
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    print("=" * 60)
+    print("  OpenEnv Labeling QA - Dataset Loader Smoke Test")
+    print("=" * 60)
+
+    # -- Task 1 ---------------------------------------------------------------
+    try:
+        t1 = load_task1()
+        print(f"\n[Sample] Task 1 ({len(t1)} total):")
+        print(f"   {t1[0]}\n")
+    except Exception as e:
+        print(f"\n[ERROR] Task 1: {e}\n")
+
+    # -- Task 2 ---------------------------------------------------------------
+    try:
+        t2 = load_task2()
+        print(f"[Sample] Task 2 ({len(t2)} total):")
+        print(f"   {t2[0]}\n")
+    except Exception as e:
+        print(f"\n[ERROR] Task 2: {e}\n")
+
+    # -- Task 3 ---------------------------------------------------------------
+    try:
+        t3 = load_task3()
+        print(f"[Sample] Task 3 ({len(t3)} total):")
+        print(f"   {t3[0]}\n")
+    except Exception as e:
+        print(f"\n[ERROR] Task 3: {e}\n")
+
+    print("=" * 60)
+    print("  Done.")
+    print("=" * 60)

environment/env.py

@@ -0,0 +1,157 @@
+import sys
+import os
+import random
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from environment.models import Observation, Action, Reward, StepResult, EpisodeState
+from environment.grader import grade
+from data.loader import load_task1, load_task2, load_task3
+from environment.noise import inject_noise_task1, inject_noise_task2, inject_noise_task3
+
+class LabelingQAEnv:
+    def __init__(self, task: str = "easy", episode_length: int = 10):
+        if task not in ["easy", "medium", "hard"]:
+            raise ValueError(f"Task must be one of 'easy', 'medium', 'hard'. Got {task}")
+            
+        self.task = task
+        self.episode_length = episode_length
+        
+        if task == "easy":
+            data = load_task1()
+            self.examples = inject_noise_task1(data)
+        elif task == "medium":
+            data = load_task2()
+            self.examples = inject_noise_task2(data)
+        elif task == "hard":
+            data = load_task3()
+            self.examples = inject_noise_task3(data)
+            
+        self.current_step = 0
+        self.cumulative_score = 0.0
+        self.done = False
+        self.current_examples = []
+        
+    def _build_observation(self, example: dict) -> Observation:
+        if self.task == "easy":
+            input_dict = {
+                "text1": example.get("text1", ""),
+                "text2": example.get("text2", "")
+            }
+            label_options = ["0", "1"]
+        elif self.task == "medium":
+            input_dict = {
+                "premise": example.get("premise", ""),
+                "hypothesis": example.get("hypothesis", "")
+            }
+            label_options = ["entailment", "neutral", "contradiction"]
+        else: # hard
+            text = example.get("text", "")
+            input_dict = {"text": text[:500]}
+            label_options = [str(i) for i in range(14)]
+            
+        return Observation(
+            example_id=str(example["id"]),
+            task=self.task,
+            input=input_dict,
+            ai_label=str(example["ai_label"]),
+            label_options=label_options,
+            episode_step=self.current_step,
+            total_steps=self.episode_length
+        )
+
+    def reset(self) -> Observation:
+        if self.episode_length > len(self.examples):
+            raise ValueError("Episode length exceeds available examples.")
+            
+        self.current_examples = random.sample(self.examples, self.episode_length)
+        self.current_step = 0
+        self.cumulative_score = 0.0
+        self.done = False
+        
+        return self._build_observation(self.current_examples[0])
+
+    def step(self, action: Action) -> StepResult:
+        if self.done:
+            raise RuntimeError("Episode is already done.")
+            
+        current_example = self.current_examples[self.current_step]
+        
+        # Pydantic v2 compatible dict dump
+        action_dict = action.model_dump() if hasattr(action, 'model_dump') else action.dict()
+        score_dict = grade(self.task, action_dict, current_example)
+        
+        reward = Reward(
+            example_id=str(current_example["id"]),
+            score=score_dict["score"],
+            reason=score_dict["reason"],
+            gold_label=str(score_dict["gold_label"])
+        )
+        
+        self.current_step += 1
+        self.cumulative_score += reward.score
+        
+        if self.current_step >= self.episode_length:
+            self.done = True
+            
+        obs = None if self.done else self._build_observation(self.current_examples[self.current_step])
+        info = {
+            "cumulative_score": self.cumulative_score,
+            "step": self.current_step
+        }
+        
+        return StepResult(
+            observation=obs,
+            reward=reward,
+            done=self.done,
+            info=info
+        )
+
+    def state(self) -> EpisodeState:
+        return EpisodeState(
+            task=self.task,
+            current_step=self.current_step,
+            total_steps=self.episode_length,
+            cumulative_score=self.cumulative_score,
+            done=self.done
+        )
+
+
+if __name__ == "__main__":
+    import sys
+    import os
+    # Dynamically inject root path specifically inside __main__ execution for simplicity
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    
+    tasks = ["easy", "medium", "hard"]
+    for t in tasks:
+        print(f"\n{'='*50}")
+        print(f"Testing LabelingQAEnv(task='{t}')")
+        print(f"{'='*50}")
+        
+        # Keep episode length short for testing
+        env = LabelingQAEnv(task=t, episode_length=3)
+        obs = env.reset()
+        
+        for i in range(3):
+            print(f"\n--- Step {i+1} ---")
+            
+            # Safely check for model_dump or standard dict wrapper
+            obs_dict = obs.model_dump() if hasattr(obs, 'model_dump') else obs.dict()
+            print(f"Observation: {obs_dict}")
+            
+            first_option = obs.label_options[0]
+            action = Action(
+                example_id=obs.example_id,
+                verdict="wrong",
+                proposed_label=first_option,
+                confidence=0.8
+            )
+            
+            res = env.step(action)
+            reward_dict = res.reward.model_dump() if hasattr(res.reward, 'model_dump') else res.reward.dict()
+            
+            print(f"Reward: {reward_dict}")
+            print(f"Cumulative Score: {env.cumulative_score}")
+            
+            obs = res.observation

environment/grader.py

@@ -0,0 +1,136 @@
+def grade_task1(action: dict, example: dict) -> dict:
+    verdict = action.get("verdict")
+    proposed_label = action.get("proposed_label")
+    gold_label_str = str(example["gold_label"])
+    is_noisy = example.get("is_noisy", False)
+    
+    score = 0.0
+    reason = ""
+    
+    if verdict == "ambiguous":
+        score = 0.3
+        reason = "Partial credit for ambiguous (task 1 has no truly ambiguous cases)."
+    elif is_noisy and verdict == "wrong":
+        score = 1.0
+        reason = "Correctly identified noisy label."
+        if proposed_label == gold_label_str:
+            score = min(1.0, score + 0.5)
+            reason = "Correctly identified noisy label and proposed the correct label."
+    elif is_noisy and verdict == "correct":
+        score = 0.0
+        reason = "Failed to identify noisy label."
+    elif not is_noisy and verdict == "correct":
+        score = 1.0
+        reason = "Correctly accepted a valid label."
+    elif not is_noisy and verdict == "wrong":
+        score = -0.5
+        reason = "Incorrectly rejected a valid label."
+        
+    return {"score": float(score), "reason": reason, "gold_label": gold_label_str}
+
+def grade_task2(action: dict, example: dict) -> dict:
+    verdict = action.get("verdict")
+    proposed_label = action.get("proposed_label")
+    gold_label_str = str(example["gold_label"])
+    is_noisy = example.get("is_noisy", False)
+    
+    score = 0.0
+    reason = ""
+    
+    if verdict == "ambiguous":
+        if gold_label_str == "neutral":
+            score = 0.5
+            reason = "Correctly identified neutral/ambiguous case."
+        else:
+            score = -0.3
+            reason = "Incorrectly labeled non-neutral case as ambiguous."
+    elif is_noisy and verdict == "wrong":
+        score = 1.0
+        reason = "Correctly identified noisy label."
+        if proposed_label == gold_label_str:
+            score = min(1.0, score + 0.5)
+            reason = "Correctly identified noisy label and proposed the correct label."
+    elif is_noisy and verdict == "correct":
+        score = 0.0
+        reason = "Failed to identify noisy label."
+    elif not is_noisy and verdict == "correct":
+        score = 1.0
+        reason = "Correctly accepted a valid label."
+    elif not is_noisy and verdict == "wrong":
+        score = -0.5
+        reason = "Incorrectly rejected a valid label."
+        
+    return {"score": float(score), "reason": reason, "gold_label": gold_label_str}
+
+def grade_task3(action: dict, example: dict) -> dict:
+    verdict = action.get("verdict")
+    proposed_label = action.get("proposed_label")
+    confidence = action.get("confidence", 1.0)
+    gold_label_str = str(example["gold_label"])
+    is_noisy = example.get("is_noisy", False)
+    
+    score = 0.0
+    reason = ""
+    
+    if verdict == "ambiguous":
+        score = 0.4
+        reason = "Partial credit for ambiguous (legal categories genuinely overlap)."
+    elif is_noisy and verdict == "wrong":
+        score = 1.0
+        reason = "Correctly identified noisy label."
+        if proposed_label == gold_label_str:
+            score = min(1.0, score + 0.5)
+            reason = "Correctly identified noisy label and proposed the correct label."
+    elif is_noisy and verdict == "correct":
+        if confidence > 0.8:
+            score = -0.5
+            reason = "Failed to identify noisy label and penalized for overconfidence."
+        else:
+            score = 0.0
+            reason = "Failed to identify noisy label."
+    elif not is_noisy and verdict == "correct":
+        score = 1.0
+        reason = "Correctly accepted a valid label."
+    elif not is_noisy and verdict == "wrong":
+        score = -0.5
+        reason = "Incorrectly rejected a valid label."
+        
+    return {"score": float(score), "reason": reason, "gold_label": gold_label_str}
+
+def grade(task: str, action: dict, example: dict) -> dict:
+    """
+    Unified entry point - routes to correct grader based on task ("easy", "medium", "hard").
+    """
+    if task == "easy":
+        return grade_task1(action, example)
+    elif task == "medium":
+        return grade_task2(action, example)
+    elif task == "hard":
+        return grade_task3(action, example)
+    else:
+        raise ValueError(f"Unknown task: {task}")
+
+if __name__ == "__main__":
+    print("--- Task 1 (Easy) Sanity Check ---")
+    ex1 = {"gold_label": 1, "ai_label": "0", "is_noisy": True}
+    act1 = {"verdict": "wrong", "proposed_label": "1", "confidence": 0.9}
+    res1 = grade("easy", act1, ex1)
+    print(f"Action:  {act1}")
+    print(f"Example: {ex1}")
+    print(f"Result:  {res1}\n")
+    
+    print("--- Task 2 (Medium) Sanity Check ---")
+    ex2 = {"gold_label": "neutral", "ai_label": "entailment", "is_noisy": True}
+    act2 = {"verdict": "ambiguous", "proposed_label": None, "confidence": 0.5}
+    res2 = grade("medium", act2, ex2)
+    print(f"Action:  {act2}")
+    print(f"Example: {ex2}")
+    print(f"Result:  {res2}\n")
+    
+    print("--- Task 3 (Hard) Sanity Check ---")
+    ex3 = {"gold_label": 11, "ai_label": "10", "is_noisy": True}
+    act3 = {"verdict": "correct", "proposed_label": None, "confidence": 0.9}
+    res3 = grade("hard", act3, ex3)
+    print(f"Action:  {act3}")
+    print(f"Example: {ex3}")
+    print(f"Result:  {res3}\n")

environment/models.py

@@ -0,0 +1,47 @@
+from typing import Literal
+from pydantic import BaseModel, Field, model_validator
+
+class Observation(BaseModel):
+    """What the agent observes at each step in the environment."""
+    example_id: str
+    task: Literal["easy", "medium", "hard"]
+    input: dict
+    ai_label: str
+    label_options: list[str]
+    episode_step: int
+    total_steps: int
+
+class Action(BaseModel):
+    """The action the agent takes for a given observation."""
+    example_id: str
+    verdict: Literal["correct", "wrong", "ambiguous"]
+    proposed_label: str | None = None
+    confidence: float = Field(ge=0.0, le=1.0)
+
+    @model_validator(mode='after')
+    def check_proposed_label(self) -> 'Action':
+        if self.verdict == "wrong" and self.proposed_label is None:
+            raise ValueError('proposed_label must be provided when verdict is "wrong"')
+        return self
+
+class Reward(BaseModel):
+    """The reward given after an action is taken."""
+    example_id: str
+    score: float = Field(ge=-1.0, le=1.0)
+    reason: str
+    gold_label: str
+
+class StepResult(BaseModel):
+    """The full return value of taking a step in the environment."""
+    observation: Observation | None
+    reward: Reward
+    done: bool
+    info: dict
+
+class EpisodeState(BaseModel):
+    """The overall state of the current episode."""
+    task: str
+    current_step: int
+    total_steps: int
+    cumulative_score: float
+    done: bool

environment/noise.py

@@ -0,0 +1,131 @@
+import random
+
+def inject_noise_task1(examples: list[dict]) -> list[dict]:
+    """
+    Task 1: gold_label is "0" or "1".
+    Noise type: random flips - 20% of examples get their label flipped to the opposite.
+    """
+    random.seed(42)
+    noised_examples = []
+    
+    for ex in examples:
+        new_ex = ex.copy()
+        gold_label = str(new_ex["gold_label"])
+        
+        # 20% chance to flip
+        if random.random() < 0.20:
+            ai_label = "1" if gold_label == "0" else "0"
+            is_noisy = True
+        else:
+            ai_label = gold_label
+            is_noisy = False
+            
+        new_ex["ai_label"] = ai_label
+        new_ex["is_noisy"] = is_noisy
+        noised_examples.append(new_ex)
+        
+    return noised_examples
+
+def inject_noise_task2(examples: list[dict]) -> list[dict]:
+    """
+    Task 2: gold_label is "entailment", "neutral", or "contradiction".
+    Noise type: systematic bias - AI always predicts "neutral" when the correct label 
+    is "contradiction". Flips 15% of "entailment" to "neutral" too.
+    """
+    random.seed(42)
+    noised_examples = []
+    
+    for ex in examples:
+        new_ex = ex.copy()
+        gold_label = str(new_ex["gold_label"])
+        
+        is_noisy = False
+        ai_label = gold_label
+        
+        if gold_label == "contradiction":
+            ai_label = "neutral"
+            is_noisy = True
+        elif gold_label == "entailment":
+            if random.random() < 0.15:
+                ai_label = "neutral"
+                is_noisy = True
+                
+        new_ex["ai_label"] = ai_label
+        new_ex["is_noisy"] = is_noisy
+        noised_examples.append(new_ex)
+        
+    return noised_examples
+
+def inject_noise_task3(examples: list[dict]) -> list[dict]:
+    """
+    Task 3: gold_label is int 0-13.
+    Noise type: confident wrong labels on edge cases - for examples where gold_label 
+    is in [0, 1, 2, 3], 30% chance of being mislabeled to a nearby category (+1 or -1). 
+    For all others, 10% random flip to any other label.
+    """
+    random.seed(42)
+    noised_examples = []
+    
+    for ex in examples:
+        new_ex = ex.copy()
+        gold_label = int(new_ex["gold_label"])
+        
+        is_noisy = False
+        ai_label = gold_label
+        
+        if gold_label in [0, 1, 2, 3]:
+            if random.random() < 0.30:
+                is_noisy = True
+                offset = random.choice([-1, 1])
+                ai_label = max(0, min(13, gold_label + offset))
+        else:
+            if random.random() < 0.10:
+                is_noisy = True
+                possible_labels = [l for l in range(14) if l != gold_label]
+                ai_label = random.choice(possible_labels)
+                
+        new_ex["ai_label"] = str(ai_label)
+        new_ex["is_noisy"] = is_noisy
+        noised_examples.append(new_ex)
+        
+    return noised_examples
+
+if __name__ == "__main__":
+    import sys
+    import os
+    sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    
+    from data.loader import load_task1, load_task2, load_task3
+    
+    # Load
+    print("Loading datasets...")
+    data1 = load_task1()
+    data2 = load_task2()
+    data3 = load_task3()
+    
+    # Task 1
+    print(f"\nTask 1 loaded: {len(data1)} examples")
+    noised1 = inject_noise_task1(data1)
+    num_noisy1 = sum(1 for ex in noised1 if ex["is_noisy"])
+    print(f"Task 1 noised: {num_noisy1} / {len(noised1)}")
+    noisy_example1 = next((ex for ex in noised1 if ex["is_noisy"]), None)
+    if noisy_example1:
+        print(f"Task 1 example: {noisy_example1}")
+        
+    # Task 2
+    print(f"\nTask 2 loaded: {len(data2)} examples")
+    noised2 = inject_noise_task2(data2)
+    num_noisy2 = sum(1 for ex in noised2 if ex["is_noisy"])
+    print(f"Task 2 noised: {num_noisy2} / {len(noised2)}")
+    noisy_example2 = next((ex for ex in noised2 if ex["is_noisy"]), None)
+    if noisy_example2:
+        print(f"Task 2 example: {noisy_example2}")
+        
+    # Task 3
+    print(f"\nTask 3 loaded: {len(data3)} examples")
+    noised3 = inject_noise_task3(data3)
+    num_noisy3 = sum(1 for ex in noised3 if ex["is_noisy"])
+    print(f"Task 3 noised: {num_noisy3} / {len(noised3)}")
+    noisy_example3 = next((ex for ex in noised3 if ex["is_noisy"]), None)
+    if noisy_example3:
+        print(f"Task 3 example: {noisy_example3}")

inference.py

@@ -0,0 +1,153 @@
+import os
+import sys
+import time
+import json
+import requests
+from openai import OpenAI
+
+API_BASE_URL = os.getenv("API_BASE_URL", "https://api.groq.com/openai/v1")
+HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("API_KEY")
+MODEL_NAME = os.getenv("MODEL_NAME", "llama-3.1-8b-instant")
+
+def build_prompt(task_name, obs):
+    base_instruction = (
+        "You are an expert data labeling quality assurance AI. "
+        "Your job is to audit the provided 'ai_label' against the input and the 'label_options'. "
+        "Respond ONLY with a valid JSON object. Do not include any other text, reasoning, or markdown formatting.\n"
+        "Required JSON keys:\n"
+        "- 'verdict': string, one of ['correct', 'wrong', 'ambiguous'].\n"
+        "- 'proposed_label': string, the correct label if verdict is 'wrong', or null if correct/ambiguous.\n"
+        "- 'confidence': float, between 0.0 and 1.0.\n\n"
+    )
+
+    input_data = obs.get("input", {})
+    ai_label = obs.get("ai_label")
+    label_options = obs.get("label_options", [])
+
+    if task_name == "easy":
+        text1 = input_data.get("sentence1", input_data.get("text1", ""))
+        text2 = input_data.get("sentence2", input_data.get("text2", ""))
+        prompt = f"Task: Easy - Binary similarity.\nText 1: {text1}\nText 2: {text2}\nAI Label: {ai_label}\nLabel Options: {label_options}"
+    elif task_name == "medium":
+        premise = input_data.get("premise", "")
+        hypothesis = input_data.get("hypothesis", "")
+        prompt = f"Task: Medium - NLI.\nPremise: {premise}\nHypothesis: {hypothesis}\nAI Label: {ai_label}\nLabel Options: {label_options}"
+    elif task_name == "hard":
+        text = input_data.get("text", "")[:300]
+        prompt = f"Task: Hard - SCOTUS legal issue classification.\nText (truncated to 300 chars): {text}...\nAI Label: {ai_label}\nLabel Options: {label_options}"
+    else:
+        prompt = f"Task: {task_name}\nInput: {input_data}\nAI Label: {ai_label}\nLabel Options: {label_options}"
+
+    return base_instruction + prompt
+
+def run_baseline() -> dict:
+    client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)
+    tasks = ["easy", "medium", "hard"]
+    
+    global_start_time = time.time()
+    
+    scores = {}
+    
+    for task in tasks:
+        session_id = f"inference_{task}"
+        
+        try:
+            start_res = requests.post(
+                "http://localhost:7860/reset", 
+                json={"task": task, "episode_length": 10, "session_id": session_id}
+            )
+            start_res.raise_for_status()
+            env_state = start_res.json()
+        except Exception as e:
+            print(f"Error resetting environment for task {task}: {e}")
+            scores[task] = 0.0
+            continue
+            
+        obs = env_state.get("observation", {})
+        done = env_state.get("done", False)
+        
+        while not done:
+            if time.time() - global_start_time > 18 * 60:
+                print("Time limit exceeding 18 minutes. Breaking early.")
+                break
+                
+            prompt = build_prompt(task, obs)
+            
+            try:
+                response = client.chat.completions.create(
+                    model=MODEL_NAME,
+                    messages=[{"role": "user", "content": prompt}],
+                    temperature=0.0
+                )
+                
+                raw_content = response.choices[0].message.content.strip()
+                
+                # Strip markdown fences if present
+                if raw_content.startswith("```json"):
+                    raw_content = raw_content[7:]
+                elif raw_content.startswith("```"):
+                    raw_content = raw_content[3:]
+                if raw_content.endswith("```"):
+                    raw_content = raw_content[:-3]
+                    
+                raw_content = raw_content.strip()
+                parsed_action = json.loads(raw_content)
+                
+                action = {
+                    "example_id": obs.get("example_id"),
+                    "verdict": parsed_action.get("verdict", "ambiguous"),
+                    "proposed_label": str(parsed_action.get("proposed_label")) if parsed_action.get("proposed_label") is not None else None,
+                    "confidence": float(parsed_action.get("confidence", 0.5))
+                }
+            except Exception as e:
+                action = {
+                    "example_id": obs.get("example_id"),
+                    "verdict": "ambiguous",
+                    "proposed_label": None,
+                    "confidence": 0.5
+                }
+            
+            payload = {
+                "session_id": session_id,
+                "action": action
+            }
+            
+            try:
+                step_res = requests.post("http://localhost:7860/step", json=payload)
+                step_res.raise_for_status()
+                env_state = step_res.json()
+                
+                obs = env_state.get("observation", {})
+                done = env_state.get("done", False)
+            except Exception as e:
+                print(f"Error stepping environment: {e}")
+                break
+                
+        # Get final score
+        try:
+            grader_res = requests.post("http://localhost:7860/grader", json={"session_id": session_id})
+            grader_res.raise_for_status()
+            score = grader_res.json().get("score", 0.0)
+            scores[task] = score
+        except Exception as e:
+            print(f"Error getting score for task {task}: {e}")
+            scores[task] = 0.0
+            
+    avg_score = sum(scores.values()) / len(scores) if scores else 0.0
+    scores["average"] = avg_score
+    return scores
+
+if __name__ == "__main__":
+    if not HF_TOKEN:
+        print("Error: HF_TOKEN or API_KEY environment variable is not set.", file=sys.stderr)
+        print("Please set it to run the baseline evaluation.", file=sys.stderr)
+        sys.exit(1)
+        
+    print("Running LabelSense Baseline...")
+    results = run_baseline()
+    
+    print("\n=== LabelSense Baseline Results ===")
+    print(f"Task: Easy    | Score: {results.get('easy', 0.0):.2f}")
+    print(f"Task: Medium  | Score: {results.get('medium', 0.0):.2f}")
+    print(f"Task: Hard    | Score: {results.get('hard', 0.0):.2f}")
+    print(f"Average Score: {results.get('average', 0.0):.2f}")

main.py

@@ -0,0 +1,128 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Dict, Optional
+
+from environment.env import LabelingQAEnv
+from environment.models import Action
+
+app = FastAPI(title="LabelSense OpenEnv")
+
+envs: Dict[str, LabelingQAEnv] = {}
+
+class ResetRequest(BaseModel):
+    task: str = "easy"
+    episode_length: int = 10
+    session_id: str = "default"
+
+class StepRequest(BaseModel):
+    session_id: str = "default"
+    example_id: str
+    verdict: str
+    proposed_label: str | None = None
+    confidence: float = 0.8
+
+@app.post("/reset")
+def reset_endpoint(req: ResetRequest):
+    try:
+        env = LabelingQAEnv(task=req.task, episode_length=req.episode_length)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+        
+    envs[req.session_id] = env
+    
+    try:
+        obs = env.reset()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+        
+    return obs.model_dump() if hasattr(obs, 'model_dump') else obs.dict()
+
+@app.post("/step")
+def step_endpoint(req: StepRequest):
+    if req.session_id not in envs:
+        raise HTTPException(status_code=404, detail="Session not found")
+        
+    env = envs[req.session_id]
+    
+    try:
+        action = Action(
+            example_id=req.example_id,
+            verdict=req.verdict,
+            proposed_label=req.proposed_label,
+            confidence=req.confidence
+        )
+        res = env.step(action)
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=str(e))
+        
+    return res.model_dump() if hasattr(res, 'model_dump') else res.dict()
+
+@app.get("/state")
+def state_endpoint(session_id: str = "default"):
+    if session_id not in envs:
+        raise HTTPException(status_code=404, detail="Session not found")
+        
+    env = envs[session_id]
+    state = env.state()
+    return state.model_dump() if hasattr(state, 'model_dump') else state.dict()
+
+@app.get("/tasks")
+def tasks_endpoint():
+    schema = {
+        "session_id": "str",
+        "example_id": "str",
+        "verdict": "str",
+        "proposed_label": "str | None",
+        "confidence": "float"
+    }
+    
+    return [
+        {
+            "name": "easy",
+            "difficulty": "Easy",
+            "description": "Binary classification setup on medical subsets",
+            "action_schema": schema
+        },
+        {
+            "name": "medium",
+            "difficulty": "Medium",
+            "description": "NLI textual entailment configuration",
+            "action_schema": schema
+        },
+        {
+            "name": "hard",
+            "difficulty": "Hard",
+            "description": "Complex multi-label edge case tagging",
+            "action_schema": schema
+        }
+    ]
+
+@app.post("/grader")
+def grader_endpoint(session_id: str = "default"):
+    if session_id not in envs:
+        raise HTTPException(status_code=404, detail="Session not found")
+        
+    env = envs[session_id]
+    state = env.state()
+    
+    return {
+        "session_id": session_id,
+        "task": state.task,
+        "cumulative_score": state.cumulative_score,
+        "total_steps": state.total_steps,
+        "done": state.done
+    }
+
+@app.post("/baseline")
+def baseline_endpoint():
+    try:
+        from baseline import run_baseline
+        return run_baseline()
+    except ImportError:
+        return {"status": "baseline not yet implemented"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/")
+def health_endpoint():
+    return {"status": "ok", "service": "LabelSense OpenEnv"}

openenv.yaml

@@ -0,0 +1,91 @@
+name: labelsense-labeling-qa
+version: 1.0.0
+description: >
+  An OpenEnv environment where an AI agent audits AI-generated labels on
+  medical and legal text datasets. The agent identifies mislabeled examples,
+  flags ambiguous cases, and proposes corrections. Simulates a real MLOps
+  data quality workflow.
+
+tags:
+  - openenv
+  - data-labeling
+  - quality-assurance
+  - medical
+  - legal
+  - nlp
+
+author: VenuGopal811
+license: MIT
+
+observation_space:
+  type: object
+  fields:
+    example_id: string
+    task: string
+    input: object
+    ai_label: string
+    label_options: array
+    episode_step: integer
+    total_steps: integer
+
+action_space:
+  type: object
+  fields:
+    example_id: string
+    verdict:
+      type: string
+      enum: [correct, wrong, ambiguous]
+    proposed_label:
+      type: string
+      nullable: true
+    confidence:
+      type: float
+      min: 0.0
+      max: 1.0
+
+tasks:
+  - name: easy
+    description: Binary medical question pair similarity labeling. ~20% noisy labels.
+    difficulty: easy
+    dataset: curaihealth/medical_questions_pairs
+    label_type: binary
+    labels: ["0", "1"]
+    expected_score_range: [0.5, 1.0]
+
+  - name: medium
+    description: 3-class NLI labeling with systematic neutral bias noise.
+    difficulty: medium
+    dataset: snli
+    label_type: multiclass
+    labels: [entailment, neutral, contradiction]
+    expected_score_range: [0.2, 0.7]
+
+  - name: hard
+    description: 14-class SCOTUS legal issue area classification with confident wrong labels.
+    difficulty: hard
+    dataset: coastalcph/lex_glue (scotus)
+    label_type: multiclass
+    labels: ["0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "10", "11", "12", "13"]
+    expected_score_range: [-0.5, 0.5]
+
+endpoints:
+  reset: POST /reset
+  step: POST /step
+  state: GET /state
+  tasks: GET /tasks
+  grader: POST /grader
+  baseline: POST /baseline
+
+baseline:
+  model: llama-3.1-8b-instant
+  script: inference.py
+  scores:
+    easy: 4.30
+    medium: 2.00
+    hard: 4.90
+    average: 3.73
+
+runtime:
+  python: "3.11"
+  framework: fastapi
+  port: 7860

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+pydantic
+datasets
+groq
+openai
+requests

tasks/task1_easy.py

@@ -0,0 +1 @@
+# Task 1 - Easy

tasks/task2_medium.py

@@ -0,0 +1 @@
+# Task 2 - Medium

tasks/task3_hard.py

@@ -0,0 +1 @@
+# Task 3 - Hard