| """ |
| JSON Array Position Bias |
| Target key-value pair at varying positions in a JSON array. |
| """ |
| import json as jsonlib |
| import logging |
| import os |
| import random |
| import time |
| import uuid |
| from typing import List, Dict, Any |
|
|
| from tqdm import tqdm |
|
|
| from src.generator import generate_text |
| from src.utils import ensure_dir, save_jsonl, save_json |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def _make_json_array(n: int, target_key: str, target_value: str, target_pos: int) -> str: |
| """Generate JSON array with target KV pair at position.""" |
| entries = [] |
| for i in range(n): |
| if i == target_pos: |
| entries.append({"key": target_key, "value": target_value}) |
| else: |
| entries.append({ |
| "key": f"key_{uuid.uuid4().hex[:8]}", |
| "value": f"val_{uuid.uuid4().hex[:8]}", |
| }) |
| return jsonlib.dumps({"records": entries}, indent=2) |
|
|
|
|
| def run_json_retrieval( |
| model_name: str, |
| num_entries: int, |
| num_examples: int, |
| out_dir: str, |
| depths: List[float] = None, |
| ) -> Dict[str, Any]: |
| ensure_dir(out_dir) |
| if depths is None: |
| depths = [0.0, 0.125, 0.25, 0.375, 0.5, 0.625, 0.75, 0.875, 1.0] |
|
|
| results = {} |
| start = time.time() |
|
|
| for depth in depths: |
| logger.info(f"[JSON] Depth {depth:.1%}") |
| preds = [] |
| for _ in tqdm(range(num_examples), desc=f"JSON {depth:.1%}", leave=False): |
| target_key = f"gold_key_{uuid.uuid4().hex[:6]}" |
| target_value = f"gold_val_{uuid.uuid4().hex[:6]}" |
| pos = int(depth * (num_entries - 1)) |
| json_str = _make_json_array(num_entries, target_key, target_value, pos) |
|
|
| prompt = ( |
| f"Find the value for the key '{target_key}' in the JSON data below.\n\n" |
| f"```json\n{json_str}\n```\n\n" |
| f"Value:" |
| ) |
| ans = generate_text( |
| [{"role": "user", "content": prompt}], |
| model_name=model_name, |
| max_new_tokens=20, |
| ) |
| correct = 1.0 if target_value.lower() in ans.lower() else 0.0 |
| preds.append({ |
| "model_answer": ans, |
| "correct": correct, |
| "target_value": target_value, |
| "depth": depth, |
| }) |
|
|
| save_jsonl(os.path.join(out_dir, f"json_depth_{depth}.jsonl"), preds) |
| acc = sum(p["correct"] for p in preds) / len(preds) if preds else 0.0 |
| results[depth] = {"accuracy": acc, "predictions": preds} |
| logger.info(f"[JSON] Depth {depth:.1%}: acc={acc:.3f}") |
|
|
| summary = { |
| "experiment": "json_retrieval", |
| "num_entries": num_entries, |
| "num_examples": num_examples, |
| "depths": {str(d): results[d]["accuracy"] for d in depths}, |
| "time_minutes": (time.time() - start) / 60, |
| } |
| save_json(os.path.join(out_dir, "json_summary.json"), summary) |
| logger.info(f"[JSON] Time={(time.time()-start)/60:.1f} min") |
| return summary |
|
|
