How to evaluate the "true" context length of your LLM using RULER

git clone https://github.com/NVIDIA/RULER.git
pip install wonderwords
cd RULER/
bash data/synthetic/json/download_qa_dataset.sh
cd scripts/data
for task in cwe fwe qa_2 vt; do python prepare.py --save_dir ./datasets --benchmark synthetic --task $task --tokenizer_path cl100k_base --tokenizer_type openai --max_seq_length 4096 --model_template_type base --num_samples 25; done

import openai
import weave
from weave import EvaluationLogger
import json
import os
from pathlib import Path
from collections import Counter, defaultdict
from typing import List, Dict, Any
import time

# ---------------- config ----------------
PROJECT = "ruler_eval"
weave.init(PROJECT)

# W&B Inference client for OSS models and judge
wandb_client = openai.OpenAI(
    base_url="https://api.inference.wandb.ai/v1",
    api_key="your_wandb_api_key",
    project=PROJECT,
    default_headers={
        "OpenAI-Project": "wandb_fc/quickstart_playground"
    }
)

# OpenAI client for gpt-5-mini via Responses API
oai_client = openai.OpenAI()

MODELS = [
    "openai/gpt-oss-20b",
    "openai/gpt-oss-120b",
    "openai/gpt-5-mini",
    "openai/gpt-5-nano",
]

JUDGE_MODEL = "gpt-5"

# Local dataset configuration - maps to your generated RULER tasks
LOCAL_DATASET_DIR = "/Users/brettyoung/Desktop/dev25/tutorials/long_cntxt/RULER/scripts/data/datasets_lg"
LOCAL_TASKS = {
    "aggregation": "cwe",      # Common Words Extraction
    "retrieval": "fwe",        # Frequency Words Extraction  
    "qa": "qa_2",             # HotpotQA
    "tracing": "vt",          # Variable Tracking
}

# -------------- dataset -----------------
def load_local_jsonl(filepath: str) -> List[Dict[str, Any]]:
    """Load a JSONL file and return list of examples"""
    examples = []
    if not os.path.exists(filepath):
        print(f"Warning: File not found: {filepath}")
        return examples
    
    with open(filepath, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    examples.append(json.loads(line))
                except json.JSONDecodeError as e:
                    print(f"Error parsing line in {filepath}: {e}")
                    continue
    return examples

def load_balanced_local(n: int = 100, seed: int = 123):
    """Load balanced samples from local RULER datasets"""
    cats = list(LOCAL_TASKS.items())
    k = len(cats)
    base, rem = divmod(n, k)

    all_examples = []
    category_counts = {}
    
    for i, (category, task_name) in enumerate(cats):
        take = base + (1 if i < rem else 0)
        if take == 0:
            continue
            
        filepath = Path(LOCAL_DATASET_DIR) / task_name / "validation.jsonl"
        print(f"Loading {filepath} (take {take})...")
        
        examples = load_local_jsonl(str(filepath))
        print(f"  {task_name} size: {len(examples)}")
        
        if len(examples) < take:
            print(f"Warning: Not enough examples in {task_name} (need {take}, have {len(examples)})")
            take = len(examples)
        
        # Add category label and select examples
        selected = examples[:take]
        for ex in selected:
            ex["category"] = category
        
        all_examples.extend(selected)
        category_counts[category] = take

    if not all_examples:
        raise RuntimeError("No examples loaded. Check dataset paths.")

    print("Category counts:", category_counts)
    return all_examples

def build_prompt(ex: Dict[str, Any]) -> str:
    """Extract prompt from RULER format"""
    # RULER format has 'input' field containing the full prompt
    if "input" in ex:
        prompt = ex["input"]
        # Add answer prefix if it exists
        if "answer_prefix" in ex:
            prompt += ex["answer_prefix"]
        return prompt
    
    # Fallback for other formats
    msgs = ex.get("messages")
    if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict):
        return "\n".join(m.get("content", "") for m in msgs if m.get("content"))
    return str(ex)

def extract_targets(ex: Dict[str, Any]) -> List[str]:
    """Extract target answers from RULER format"""
    # RULER format uses 'outputs' field
    outputs = ex.get("outputs", [])
    if isinstance(outputs, list):
        return [str(x) for x in outputs if x is not None]
    elif outputs is not None:
        return [str(outputs)]
    
    # Fallback to other common fields
    for field in ["expected_answer", "answer", "target"]:
        v = ex.get(field)
        if v is not None:
            if isinstance(v, list):
                return [str(x) for x in v if x is not None]
            return [str(v)]
    
    return []

# -------------- inference ----------------
def chat_once(model: str, prompt: str) -> str | None:
    try:
        if model == "openai/gpt-5-mini":
            resp = oai_client.responses.create(
                model="gpt-5-mini",
                reasoning={"effort": "low"},
                input=[{"role": "user", "content": prompt}],
            )
            return resp.output_text

        if model == "openai/gpt-5-nano":
            resp = oai_client.responses.create(
                model="gpt-5-nano",
                reasoning={"effort": "low"},
                input=[{"role": "user", "content": prompt}],
            )
            return resp.output_text
                    
        if model in ["openai/gpt-oss-20b", "openai/gpt-oss-120b"]:
            resp = wandb_client.chat.completions.create(
                model=model,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.0,
                max_tokens=512,
            )
            return resp.choices[0].message.content
            
    except Exception as e:
        print(f"Error in chat_once for {model}: {e}")
        return None

def judge(pred: str | None, golds: List[str], ctx: str, problem_excerpt: str) -> float:
    if pred is None or not golds:
        return 0.0
    gold_text = str(golds)
    jp = (
        "You are an evaluator.\n"
        "Use the problem excerpt to infer whether order matters for this task. "
        "If the instruction implies a specific ordering or exact string match, be strict. "
        "If it asks for entities, spans, counts, or unordered sets, allow order-free matching. "
        "Minor formatting and whitespace should be ignored.\n\n"
        f"Problem excerpt (first 300 chars): {problem_excerpt}\n\n"
        f"Gold answers (one may be correct): {gold_text}\n"
        f"Model prediction: {pred}\n\n"
        "Return a single word: CORRECT if the prediction semantically matches any gold answer given the instruction, otherwise INCORRECT. Ignore minor formatting differences and like added *'s or just a more verbose answer."
    )
    try:
        # Use GPT-5 via OpenAI Responses API for judging
        resp = oai_client.responses.create(
            model=JUDGE_MODEL,
            reasoning={"effort": "low"},
            input=[{"role": "user", "content": jp}],
        )
        out = resp.output_text if resp and resp.output_text else ""
        return 1.0 if out and out.strip().upper().startswith("CORRECT") else 0.0
    except Exception as e:
        print(f"Error in judge: {e}")
        return 0.0

# -------------- eval ---------------
def run_eval_for_model(model_name: str, examples):
    safe_model_name = model_name.replace("/", "_").replace("-", "_")
    print("Using safe model name for logger:", safe_model_name)

    ev = EvaluationLogger(model=safe_model_name, dataset="RULER_Local_4K")

    per_cat_j = defaultdict(list)

    for i, ex in enumerate(examples):
        time.sleep(5)  # Rate limit safety 
        print(f"\n--- Example {i} ({ex['category']}) ---")
        print("Raw example keys:", list(ex.keys()))

        prompt = build_prompt(ex)
        print("Prompt preview:", prompt[:200], "..." if len(prompt) > 200 else "")

        golds = extract_targets(ex)
        print("Gold answers:", golds)

        pred = chat_once(model_name, prompt)
        print("Model prediction:", pred)

        pred_logger = ev.log_prediction(
            inputs={"category": ex["category"], "prompt": prompt, "gold": golds},
            output={"prediction": pred}
        )

        problem_excerpt = prompt[:300]
        jscore = judge(pred, golds, prompt, problem_excerpt)
        print("Judge score:", jscore)

        pred_logger.log_score("judge", jscore)
        pred_logger.finish()

        per_cat_j[ex["category"]].append(jscore)

        if i % 5 == 0:
            print(f"Progress checkpoint -> {i}: Judge {jscore}")

    overall_j = sum(sum(v) for v in per_cat_j.values()) / max(1, len(examples))
    per_cat_j_avg = {k: sum(v)/len(v) for k, v in per_cat_j.items()}

    print(f"\nSummary for {model_name}")
    print("Overall Judge:", overall_j)
    print("Per-category Judge:", per_cat_j_avg)

    ev.log_summary({
        "overall_judge_correct": overall_j,
        "per_category_judge_correct": per_cat_j_avg,
        "model": model_name,
    })

def main():
    print("Loading local RULER datasets...")
    examples = load_balanced_local(n=100, seed=123) 
    print("Loaded dataset with", len(examples), "examples")

    for m in MODELS:
        print("\n==============================")
        print("Evaluating model:", m)
        print("==============================")
        run_eval_for_model(m, examples)

    print("Done")

if __name__ == "__main__":
    main()