[content] tab-model-eval/src/evaluation/sglang_eval.py (python)
<<<<<<< HEAD
    
    # Single-file mode (backward compatible)
    generations_file: str = ""
    evaluations_file: str = ""
    eval_step: int = 0
    
    # Batch mode: comma-separated lists of files and steps
    # When these are provided, they take precedence over single-file args
    generations_files: str = ""  # Comma-separated list of generation files
    evaluations_files: str = ""  # Comma-separated list of evaluation output files
    eval_steps: str = ""  # Comma-separated list of eval steps (integers)
    
=======
    generations_file: str = "data/eval/handcrafted_test_cases/handcrafted_generations.jsonl"
    evaluations_file: str = "data/eval/handcrafted_test_cases/handcrafted_evaluations.jsonl"
>>>>>>> dev/lora_run
    limit: int = -1
    system_prompt_file: str = "data/prompts/judge_system_prompt_v2.md"
    judge_name: str = "default"
    judge_prompt_file: str = "data/prompts/judge_prompt_v2.md"
    judge_prompt_file_with_context: str = "data/prompts/judge_prompt_v2_with_context.md"
    eval_step: int = 0

    # Local logging for offline mode
    use_local_logger: bool = False
    local_log_dir: str = "data/eval/local_logs"

    # Server-related (sglang)
    judge_model_path: str = "Qwen/Qwen3-Coder-30B-A3B-Instruct"
    server_host: str = "0.0.0.0"
    server_port: int = 30000
    context_length: int = 40960
    problem_length: int = 40960
    api_key: str = "EMPTY"  # sglang's OpenAI-compatible server ignores this value
    mem_fraction_static: float = 0.95
    tp_size: int = 1

    # Client-related
    temperature: float = 0.7
    top_p: float = 0.8
    presence_penalty: float = 1.5
    top_k: int = 20
    min_p: float = 0.0
    enable_thinking: bool = True

    # HTTP / client config
    concurrency: int = 16
    max_connections: int = 256
    keepalive: int = 60
    max_attempts: int = 6
    timeout: float = 30.0

    # Control whether to launch server from this script
    launch_server: bool = True
    # Extra args passed to `sglang.launch_server` if needed
    extra_server_args: Optional[List[str]] = None
    
    def get_eval_jobs(self) -> List[tuple[str, str, int]]:
        """
        Returns a list of (generations_file, evaluations_file, eval_step) tuples.
        If batch mode args are provided, uses those. Otherwise falls back to single-file mode.
        """
        if self.generations_files and self.evaluations_files and self.eval_steps:
            # Batch mode
            gen_files = [f.strip() for f in self.generations_files.split(",") if f.strip()]
            eval_files = [f.strip() for f in self.evaluations_files.split(",") if f.strip()]
            steps = [int(s.strip()) for s in self.eval_steps.split(",") if s.strip()]
            
            if not (len(gen_files) == len(eval_files) == len(steps)):
                raise ValueError(
                    f"Batch mode requires equal-length lists for generations_files ({len(gen_files)}), "
                    f"evaluations_files ({len(eval_files)}), and eval_steps ({len(steps)})"
                )
            
            return list(zip(gen_files, eval_files, steps))
        elif self.generations_file and self.evaluations_file:
            # Single-file mode (backward compatible)
            return [(self.generations_file, self.evaluations_file, self.eval_step)]
        else:
            raise ValueError(
                "Either provide single-file args (generations_file, evaluations_file) "
                "or batch args (generations_files, evaluations_files, eval_steps)"
            )


# ----------------------------
# Dataset helpers
# ----------------------------
def load_dataset(filepath):
    with open(filepath, "r") as f:
        return json.loads(f.read())


def estimate_token_count(messages: List[Dict[str, str]]) -> int:
    """
    Rough estimate of token count for a list of messages.
    Assumes ~3 characters per token as a conservative estimate.
    """
    total_chars = sum(len(msg.get("content", "")) for msg in messages)
    return total_chars // 3


def filter_tasks_by_context_length(
    test_cases: List[Dict[str, Any]],
    system_prompt: str,
    prompt_template: str,
    max_context_length: int = 40960,
    problem_length: int = 40960,
    buffer_tokens: int = 512,  # Reserve space for response
    include_context: bool = False,
) -> tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
    """
    Filter out test cases whose context would exceed the model's context length.
    Returns (valid_cases, skipped_cases)
    """
    valid_cases = []
    skipped_cases = []

    for tc in test_cases:
        # Estimate tokens for system prompt + context
        messages = [{"role": "system", "content": system_prompt}]
        if include_context:
            messages.extend(tc["context"])
        messages.append({"role": "user", "content": prompt_template})
        estimated_tokens = estimate_token_count(messages)

        length = estimated_tokens + buffer_tokens
        if length <= max_context_length and length <= problem_length:
            valid_cases.append(tc)
        else:
            print(
                f"Skipping {tc['task_id']}: estimated {estimated_tokens} tokens (limit: {max_context_length}, problem_length: {problem_length})"
            )
            skipped_cases.append(
                {
                    "task_id": tc["task_id"],
                    "estimated_tokens": estimated_tokens,
                    "reason": "context_too_long",
                }
            )

    return valid_cases, skipped_cases


# ----------------------------
# Eval logic
# ----------------------------
async def evaluate_generated_command(
    client: AsyncOpenAI,
    sem: asyncio.Semaphore,
    test_case: Dict[str, Any],
    args: Args,
    system_prompt: str,
    prompt_template: str,
    include_context: bool,
) -> Dict[str, Any]:
    """
    Handles a single evaluation task with concurrency control and retries.
    """
    async with sem:
        delay = 0.25

        if test_case.get("error", None) is not None:
            print(f"Returning failure object for task {test_case['task_id']} due to error")
            return {
                "task_id": test_case["task_id"],
                "error": test_case["error"],
                "is_correct": 0,
                "average_score": 0.0,
            }

        samples = test_case.get("samples", [])
        if not samples:
            print(f"Returning failure object for task {test_case['task_id']} due to no samples")
            return {
                "task_id": test_case["task_id"],
                "error": "No samples",
                "is_correct": 0,
                "average_score": 0.0,
            }

        sample_results = []
        for sample in samples:
            for attempt in range(args.max_attempts):
                try:
                    format_dict = {
                        "expected": test_case["expected_command"],
                        "generated": sample["generated_command"],
                    }
                    if include_context:
                        format_dict["context"] = json.dumps(test_case["context"], indent=2)
                    prompt = prompt_template.format(**format_dict)

                    messages = [
                        {
                            "role": "system",
                            "content": system_prompt,
                        },
                        {"role": "user", "content": prompt},
                    ]

                    resp = await client.chat.completions.create(
                        model=args.judge_name,
                        messages=messages,
                        temperature=args.temperature,
                        top_p=args.top_p,
                        presence_penalty=args.presence_penalty,
                        response_format={"type": "json_object"},
                        extra_body={
                            "top_k": args.top_k,
                            "chat_template_kwargs": {"enable_thinking": args.enable_thinking},
                        },
                    )

                    thinking_trace = getattr(resp.choices[0].message, "reasoning_content", "")
                    result = json.loads(resp.choices[0].message.content)
                    equivalent = result.get("equivalent", 0)

                    sample_results.append(
                        {
                            "generated_command": sample["generated_command"],
                            "thinking_trace": thinking_trace,
                            "evaluation_results": result,
                            "equivalent": equivalent,
                            "exact_match": sample["exact_match"],
                        }
                    )
                    break

                except BadRequestError as e:
                    print(
                        f"Returning failure object for task {test_case['task_id']} due to BadRequestError: {e}"
                    )
                    sample_results.append(
                        {
                            "task_id": test_case["task_id"],
                            "error": str(e),
                            "equivalent": 0,
                            "exact_match": 0,
                        }
                    )
                    break

                except Exception as e:
                    print(f"Error on task {test_case['task_id']}: {e}")
                    if attempt == args.max_attempts - 1:
                        print(f"Returning failure object for task {test_case['task_id']}")
                        sample_results.append(
                            {
                                "task_id": test_case["task_id"],
                                "error": str(e),
                                "equivalent": 0,
                            }
                        )
                    await asyncio.sleep(delay)
                    delay *= 2

        # Compute avg@n and pass@n
        num_judge_matches = sum(s.get("equivalent", 0) for s in sample_results)
        judge_avg_at_n = num_judge_matches / len(sample_results)
        judge_pass_at_n = int(num_judge_matches > 0)
        num_exact_matches = test_case.get("num_exact_matches", 0)

        return {
            "task_id": test_case["task_id"],
            "context": test_case["context"],
            "expected_command": test_case["expected_command"],
            "sample_evaluations": sample_results,
            "num_samples": len(sample_results),
            "num_judge_matches": num_judge_matches,
            "judge_avg_at_n": judge_avg_at_n,
            "judge_pass_at_n": judge_pass_at_n,
            "num_exact_matches": num_exact_matches,
        }


async def run_single_eval(
    args: Args,
    generations_file: str,
    evaluations_file: str,
    eval_step: int,
    client: AsyncOpenAI,
    sem: asyncio.Semaphore,
    system_prompt: str,
    prompt_template: str,
    include_context: bool,
    logger: Optional[LocalLogger] = None,
) -> Dict[str, Any]:
    """
    Evaluate a single generations file and write results to evaluations file.
    Uses shared client and semaphore for efficiency in batch mode.
    Returns the evaluation scores dictionary.
    """
    print(f"
{'='*60}")
    print(f"Evaluating: {generations_file}")
    print(f"Output: {evaluations_file}")
    print(f"Step: {eval_step}")
    print(f"{'='*60}")
    
    loaded_data = load_dataset(generations_file)
    test_cases = loaded_data["generation_results"]

    config_generations = loaded_data["config_generations"]
    config_evaluations = args.__dict__
    metadata = {
        "config_generations": config_generations,
        "config_evaluations": config_evaluations,
    }

<<<<<<< HEAD
=======
    wandb_init_kwargs = {
        "project": args.wandb_project,
        "name": args.wandb_name,
        "tags": args.wandb_tags,
        "group": args.wandb_group,
        "config": metadata,
    }

    if args.wandb_id:
        wandb_init_kwargs.update(
            {
                "id": args.wandb_id,
                "resume": "allow",
            }
        )
    wandb.init(**wandb_init_kwargs)

>>>>>>> dev/lora_run
    if args.limit > 0:
        test_cases = test_cases[: args.limit]

    # Filter out tasks with context that's too long
    test_cases, skipped_cases = filter_tasks_by_context_length(
        test_cases,
        system_prompt=system_prompt,
        prompt_template=prompt_template,
        max_context_length=args.context_length,
        problem_length=args.problem_length,
        buffer_tokens=512,
        include_context=include_context,
    )

    print(f"
Filtered dataset:")
    print(f"  Valid test cases: {len(test_cases)}")
    print(f"  Skipped (too long): {len(skipped_cases)}")
    print()

    # Clean output
    if os.path.exists(evaluations_file):
        os.remove(evaluations_file)

    tasks = [
        evaluate_generated_command(
            client, sem, tc, args, system_prompt, prompt_template, include_context
        )
        for tc in test_cases
    ]

    print(f"Running {len(test_cases)} test cases with concurrency={args.concurrency} ...")
    results: List[Dict[str, Any]] = []

    # progress bar over async tasks
    for coro in tqdm_asyncio.as_completed(tasks, total=len(tasks)):
        results.append(await coro)

    # sort the results by task_id
    results.sort(key=lambda x: x["task_id"])

    os.makedirs(os.path.dirname(evaluations_file), exist_ok=True)
    total_judge_avg_at_n = sum(r.get("judge_avg_at_n", 0) for r in results) / len(results)
    total_judge_pass_at_n = sum(r.get("judge_pass_at_n", 0) for r in results)

    total_exact_match_avg_at_n = loaded_data["generation_scores"]["total_exact_match_avg_at_n"]
    total_exact_match_pass_at_n = loaded_data["generation_scores"]["total_exact_match_pass_at_n"]

<<<<<<< HEAD
    # Prepare metrics to log
    metrics_to_log = {
        "eval_step": eval_step,
        f"{args.wandb_eval_type}/total_test_cases": len(test_cases),
        f"{args.wandb_eval_type}/num_samples_per_task": loaded_data["config_generations"][
            "num_samples"
        ],
        f"{args.wandb_eval_type}/total_judge_avg_at_n": total_judge_avg_at_n,
        f"{args.wandb_eval_type}/total_judge_pass_at_n": total_judge_pass_at_n,
        f"{args.wandb_eval_type}/total_exact_match_avg_at_n": total_exact_match_avg_at_n,
        f"{args.wandb_eval_type}/total_exact_match_pass_at_n": total_exact_match_pass_at_n,
    }
    
    # Log metrics using appropriate logger
    if args.use_local_logger and logger is not None:
        logger.log(metrics_to_log)
    else:
        wandb.log(metrics_to_log)

    with open(evaluations_file, "w") as f:
=======
    wandb.log(
        {
            f"eval_step": args.eval_step, # global eval_step since it is the step number of the checkpoint
            f"{args.wandb_eval_type}/total_test_cases": len(test_cases),
            f"{args.wandb_eval_type}/num_samples_per_task": loaded_data["config_generations"][
                "num_samples"
            ],
            f"{args.wandb_eval_type}/total_judge_avg_at_n": total_judge_avg_at_n,
            f"{args.wandb_eval_type}/total_judge_pass_at_n": total_judge_pass_at_n,
            f"{args.wandb_eval_type}/total_exact_match_avg_at_n": total_exact_match_avg_at_n,
            f"{args.wandb_eval_type}/total_exact_match_pass_at_n": total_exact_match_pass_at_n,
        }
    )

    with open(args.evaluations_file, "w") as f:
>>>>>>> dev/lora_run