Enh(analysis): Add script to pick instances for llm analysis

Author: klieret

codeclash/analysis/llm_as_judge/big_questions.py

@@ -17,7 +17,7 @@
 from rich.table import Table
 from typing_extensions import Any
 
-from codeclash.analysis.llm_as_judge.utils import FileLock, Instance, get_instances
+from codeclash.analysis.llm_as_judge.utils import FileLock, Instance, InstanceBatch, get_instances
 from codeclash.utils.log import get_logger
 
 logger = get_logger("BigQuestionsEvaluator", emoji="🤖")
@@ -204,17 +204,28 @@ def evaluate_bulk(self, instances: list[Instance], *, n_workers: int = 1) -> Non
                         )
 
 
+def load_instances_from_path(path: Path) -> list[Instance]:
+    if path.is_file() and path.suffix == ".json":
+        logger.info(f"Loading instances from batch file: {path}")
+        batch = InstanceBatch.model_validate_json(path.read_text())
+        return batch.instances
+    logger.info(f"Loading instances from directory: {path}")
+    return get_instances(path)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("input_dir", type=Path, nargs="+", help="Path to the input dir(s)")
+    parser.add_argument(
+        "input_dir", type=Path, nargs="+", help="Path to the input dir(s) or to instance batch json files"
+    )
     parser.add_argument("--shuffle", action="store_true", help="Shuffle instances before processing")
     parser.add_argument("-n", "--n-workers", type=int, default=1, help="Number of parallel workers (default: 1)")
     args = parser.parse_args()
 
     config = BigQuestionsConfig.model_validate(yaml.safe_load(config_path.read_text()))
     instances = []
-    for input_dir in args.input_dir:
-        instances.extend(get_instances(input_dir))
+    for input_path in args.input_dir:
+        instances.extend(load_instances_from_path(input_path))
     big_questions = BigQuestions(config)
     if args.shuffle:
         random.shuffle(instances)

codeclash/analysis/llm_as_judge/get_instances.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python3
+
+"""Get instances based on specific pattern."""
+
+import argparse
+from pathlib import Path
+
+from codeclash.analysis.llm_as_judge.utils import InstanceBatch, get_instances
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input_dir", type=Path, help="Path to the input log dir")
+    parser.add_argument("-o", "--output-file", type=Path, help="Path to the output file", default="instances.json")
+    args = parser.parse_args()
+
+    SELECTED_ROUNDS = [1, 5, 10, 15]
+    SELECTED_GAME = "BattleSnake"
+
+    # first get all instances, then filter
+    instances = sorted(get_instances(args.input_dir), key=lambda x: x.instance_id)
+    print(f"Found {len(instances)} instances")
+    instances = [instance for instance in instances if SELECTED_GAME in instance.tournament_name]
+    print(f"Filtered to {len(instances)} instances because of game name")
+    instances = [instance for instance in instances if instance.round_number in SELECTED_ROUNDS]
+    print(f"Filtered to {len(instances)} instances because of round number")
+
+    # OK, now it gest more complicated, because we only want to keep one
+    # instance per combination of game, lm, and round
+    unique = {}
+    for instance in instances:
+        key = (instance.tournament_name, instance.get_lm_name(), instance.round_number)
+        if key not in unique:
+            unique[key] = instance
+    instances = list(unique.values())
+    print(f"Filtered to {len(instances)} instances after deduplication for game repetitions")
+
+    batch = InstanceBatch(instances=instances)
+    args.output_file.write_text(batch.model_dump_json())
+    print(f"Wrote {len(instances)} instances to {args.output_file}")

codeclash/analysis/llm_as_judge/utils.py

@@ -1,4 +1,5 @@
 import fcntl
+import json
 from pathlib import Path
 
 from pydantic import BaseModel
@@ -31,6 +32,15 @@ class Instance(BaseModel):
     def instance_id(self) -> str:
         return f"{self.tournament_name}__{self.player_name}__r{self.round_number}"
 
+    def get_lm_name(self) -> str:
+        metadata_path = self.trajectory_path.parent.parent.parent / "metadata.json"
+        metadata = json.loads(metadata_path.read_text())
+        return metadata["config"]["players"][self.player_name]["config"]["model"]["model_name"]
+
+
+class InstanceBatch(BaseModel):
+    instances: list[Instance]
+
 
 def find_tournament_folders(input_dir: Path) -> list[Path]:
     return [d.parent for d in input_dir.rglob("metadata.json")]