run_llm_eval_opensource.py

#!/usr/bin/env python3
"""
pass@1 assessment using open-source code model (reproducible).

Uses a small open-source code generation model (Salesforce/codegen-350M-mono)
instead of MiniMax M2.5 for reproducible assessment.

Compares: Full Context vs CTX Adaptive Trigger
Metric: pass@1 (similarity-based judgment of generated code vs reference)
"""

import json
import os
import re
import sys
import time
from datetime import datetime
from typing import Dict, List, Optional, Tuple

import numpy as np
import torch
from scipy import stats

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from src.evaluator.llm_quality import (
    FunctionSample,
    sample_functions,
    build_full_context,
    build_adaptive_context,
    extract_functions_from_file,
)
from src.retrieval.full_context import estimate_tokens


class OpenSourceCodeAssessor:
    """Code generation assessor using open-source models."""

    def __init__(self, model_name: str = "Salesforce/codegen-350M-mono", device: str = None):
        from transformers import AutoTokenizer, AutoModelForCausalLM

        self.model_name = model_name
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        print(f"  Loading model: {model_name}")
        print(f"  Device: {self.device}")

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        print(f"  Tokenizer loaded, loading model weights...", flush=True)

        # Use CPU for 350M model (avoids CUDA assert issues with codegen tokenizer)
        # 350M params on CPU with fp32 is fast enough (~2s per inference)
        if "350M" in model_name or "350m" in model_name:
            self.device = "cpu"
            print(f"  Using CPU for small model (avoids CUDA tokenizer issues)", flush=True)

        dtype = torch.float16 if self.device == "cuda" else torch.float32
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            dtype=dtype,
        )
        print(f"  Moving model to {self.device}...", flush=True)
        self.model = self.model.to(self.device)
        print(f"  Model on device.", flush=True)

        # Set pad token if not set
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        param_count = sum(p.numel() for p in self.model.parameters()) / 1e6
        print(f"  Model loaded ({param_count:.0f}M params)")

    def generate_code(self, task_prompt: str, context: str, max_new_tokens: int = 128) -> str:
        """Generate code given a task prompt and context."""
        # Keep context short to fit in model's 2048 token context window
        # Reserve ~256 tokens for generation + ~200 for task prompt
        max_context_chars = 1200 * 4  # ~1200 tokens of context
        truncated_context = context[:max_context_chars]

        prompt = (
            f"# Context:\n{truncated_context}\n\n"
            f"# Task: {task_prompt}\n"
        )

        inputs = self.tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1600,  # Leave 448 tokens for generation
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=0.2,
                top_p=0.95,
                do_sample=True,
                pad_token_id=self.tokenizer.eos_token_id,
            )

        # Decode only the generated tokens
        generated = self.tokenizer.decode(
            outputs[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        )

        return generated.strip()

    def assess_similarity(self, generated: str, reference: str) -> Tuple[bool, float, str]:
        """Assess if generated code is functionally similar to reference.

        Uses a combination of:
        1. Exact/substring match
        2. Token overlap (Jaccard similarity)
        3. Coverage (fraction of reference tokens present)

        Returns:
            Tuple of (passed, similarity_score, assessment_detail)
        """
        # Normalize whitespace for comparison
        gen_norm = re.sub(r'\s+', ' ', generated.strip())
        ref_norm = re.sub(r'\s+', ' ', reference.strip())

        # Check exact match (after normalization)
        if gen_norm == ref_norm:
            return True, 1.0, "exact_match"

        # Extract key tokens (identifiers, keywords, operators)
        def extract_tokens(code: str) -> set:
            tokens = set(re.findall(r'[a-zA-Z_]\w+|[+\-*/=<>!&|%]+|\d+', code))
            # Remove very common tokens
            common = {'self', 'return', 'def', 'class', 'if', 'else', 'for', 'in',
                      'while', 'try', 'except', 'import', 'from', 'and', 'or', 'not',
                      'None', 'True', 'False', 'pass', 'the', 'is', 'as', 'with'}
            return tokens - common

        gen_tokens = extract_tokens(generated)
        ref_tokens = extract_tokens(reference)

        if not ref_tokens:
            return False, 0.0, "empty_reference"

        # Jaccard similarity
        intersection = gen_tokens & ref_tokens
        union = gen_tokens | ref_tokens
        jaccard = len(intersection) / len(union) if union else 0.0

        # Coverage: what fraction of reference tokens appear in generated
        coverage = len(intersection) / len(ref_tokens) if ref_tokens else 0.0

        # Combined score
        similarity = 0.6 * coverage + 0.4 * jaccard

        # Threshold for "pass"
        passed = similarity >= 0.35
        detail = f"coverage={coverage:.3f}, jaccard={jaccard:.3f}, combined={similarity:.3f}"

        return passed, similarity, detail


def compute_wilson_ci(passed: int, total: int, confidence: float = 0.95) -> Tuple[float, float]:
    """Compute Wilson score confidence interval."""
    if total == 0:
        return (0.0, 0.0)

    p = passed / total
    z = stats.norm.ppf(1 - (1 - confidence) / 2)

    denominator = 1 + z**2 / total
    center = (p + z**2 / (2 * total)) / denominator
    spread = z * np.sqrt((p * (1 - p) + z**2 / (4 * total)) / total) / denominator

    return (max(0.0, center - spread), min(1.0, center + spread))


def compute_mcnemar_test(per_sample: list) -> dict:
    """Compute McNemar test for paired comparison."""
    a, b, c, d = 0, 0, 0, 0

    for s in per_sample:
        fc = s["strategies"].get("full_context", {}).get("passed", False)
        at = s["strategies"].get("adaptive_trigger", {}).get("passed", False)

        if fc and at:
            a += 1
        elif fc and not at:
            b += 1
        elif not fc and at:
            c += 1
        else:
            d += 1

    if b + c == 0:
        p_value = 1.0
        chi2 = 0.0
    else:
        chi2 = (abs(b - c) - 1) ** 2 / (b + c)
        p_value = 1.0 - stats.chi2.cdf(chi2, df=1)

    return {
        "contingency": {"both_pass": a, "fc_only": b, "at_only": c, "both_fail": d},
        "chi2": chi2,
        "p_value": p_value,
    }


def run_opensource_assessment(
    project_path: str,
    n_samples: int = 30,
    model_name: str = "Salesforce/codegen-350M-mono",
    results_dir: str = "benchmarks/results",
    seed: int = 42,
) -> dict:
    """Run pass@1 assessment using open-source code model.

    Args:
        project_path: Path to Python project
        n_samples: Number of function samples
        model_name: HuggingFace model name
        results_dir: Output directory
        seed: Random seed

    Returns:
        Results dictionary
    """
    os.makedirs(results_dir, exist_ok=True)
    timestamp = datetime.now().isoformat()
    project_name = os.path.basename(project_path)

    print("=" * 60)
    print("CTX Open-Source LLM Quality Assessment (pass@1)")
    print("=" * 60)
    print(f"  Model: {model_name}")
    print(f"  Project: {project_name} ({project_path})")
    print(f"  Samples: {n_samples}")
    print(f"  Seed: {seed}")
    print()

    # Step 1: Load model
    print("[1/4] Loading model...")
    try:
        assessor = OpenSourceCodeAssessor(model_name=model_name)
    except Exception as e:
        error_msg = f"Failed to load model {model_name}: {e}"
        print(f"  ERROR: {error_msg}")

        # Try fallback model
        fallback = "Salesforce/codegen-350M-mono"
        if model_name != fallback:
            print(f"  Trying fallback model: {fallback}")
            try:
                assessor = OpenSourceCodeAssessor(model_name=fallback)
                model_name = fallback
            except Exception as e2:
                error_msg += f"\n  Fallback also failed: {e2}"
                print(f"  ERROR: {error_msg}")
                _save_failure_report(error_msg, results_dir, timestamp)
                return {"error": error_msg}
    print()

    # Step 2: Sample functions
    print("[2/4] Sampling functions...")
    functions = sample_functions(project_path, n=n_samples, seed=seed)
    print(f"  Sampled {len(functions)} functions")
    for i, func in enumerate(functions[:5]):
        print(f"    {i+1}. {func.name} ({func.file_path})")
    if len(functions) > 5:
        print(f"    ... and {len(functions) - 5} more")
    print()

    # Step 3: Build full context
    print("[3/4] Building contexts...")
    full_context = build_full_context(project_path, max_tokens=1200)
    full_ctx_tokens = estimate_tokens(full_context)
    print(f"  Full context: {full_ctx_tokens} tokens")
    print()

    # Step 4: Run assessment
    print(f"[4/4] Running assessment on {len(functions)} functions...")
    all_per_sample = []
    start_time = time.time()

    for i, func in enumerate(functions):
        print(f"\n  [{i+1}/{len(functions)}] {func.name} ({func.file_path})")

        task_prompt = f"Implement function '{func.name}'"
        if func.docstring:
            task_prompt += f" with this specification: {func.docstring[:200]}"
        task_prompt += f"\n\nFunction signature: {func.signature}"

        sample_result = {
            "index": i,
            "function_name": func.name,
            "file_path": func.file_path,
            "docstring": func.docstring[:200] if func.docstring else "",
            "reference_lines": len(func.body.splitlines()),
            "strategies": {},
        }

        # Strategy 1: Full Context
        print("    Full Context: ", end="", flush=True)
        try:
            generated_full = assessor.generate_code(task_prompt, full_context)
            passed_full, sim_full, detail_full = assessor.assess_similarity(
                generated_full, func.full_source
            )
            sample_result["strategies"]["full_context"] = {
                "passed": passed_full,
                "similarity": sim_full,
                "context_tokens": full_ctx_tokens,
                "generated_preview": generated_full[:200],
                "assessment_detail": detail_full,
                "error": None,
            }
            print(f"{'PASS' if passed_full else 'FAIL'} (sim={sim_full:.3f})")
        except Exception as e:
            sample_result["strategies"]["full_context"] = {
                "passed": False,
                "similarity": 0.0,
                "context_tokens": full_ctx_tokens,
                "error": str(e),
            }
            print(f"ERROR: {e}")

        # Strategy 2: Adaptive Trigger
        print("    Adaptive Trigger: ", end="", flush=True)
        try:
            adaptive_ctx = build_adaptive_context(project_path, func, max_tokens=800)
            adaptive_tokens = estimate_tokens(adaptive_ctx)

            generated_adaptive = assessor.generate_code(task_prompt, adaptive_ctx)
            passed_adaptive, sim_adaptive, detail_adaptive = assessor.assess_similarity(
                generated_adaptive, func.full_source
            )
            sample_result["strategies"]["adaptive_trigger"] = {
                "passed": passed_adaptive,
                "similarity": sim_adaptive,
                "context_tokens": adaptive_tokens,
                "generated_preview": generated_adaptive[:200],
                "assessment_detail": detail_adaptive,
                "error": None,
            }
            print(f"{'PASS' if passed_adaptive else 'FAIL'} (sim={sim_adaptive:.3f}, ctx={adaptive_tokens} tok)")
        except Exception as e:
            sample_result["strategies"]["adaptive_trigger"] = {
                "passed": False,
                "similarity": 0.0,
                "error": str(e),
            }
            print(f"ERROR: {e}")

        all_per_sample.append(sample_result)

        # Check time budget
        elapsed = time.time() - start_time
        if elapsed > 900 and i < len(functions) - 1:  # 15 min limit
            remaining = len(functions) - i - 1
            print(f"\n  WARNING: Time budget exceeded ({elapsed:.0f}s). "
                  f"Stopping after {i+1} samples ({remaining} remaining)")
            break

    elapsed = time.time() - start_time

    # Compute final metrics
    print(f"\n{'='*60}")
    print("Computing final metrics...")

    total_n = len(all_per_sample)

    fc_passed = sum(
        1 for s in all_per_sample
        if s["strategies"].get("full_context", {}).get("passed", False)
    )
    at_passed = sum(
        1 for s in all_per_sample
        if s["strategies"].get("adaptive_trigger", {}).get("passed", False)
    )

    fc_pass1 = fc_passed / total_n if total_n > 0 else 0.0
    at_pass1 = at_passed / total_n if total_n > 0 else 0.0

    fc_ci = compute_wilson_ci(fc_passed, total_n)
    at_ci = compute_wilson_ci(at_passed, total_n)

    fc_sims = [
        s["strategies"].get("full_context", {}).get("similarity", 0.0)
        for s in all_per_sample
    ]
    at_sims = [
        s["strategies"].get("adaptive_trigger", {}).get("similarity", 0.0)
        for s in all_per_sample
    ]

    fc_avg_sim = float(np.mean(fc_sims)) if fc_sims else 0.0
    at_avg_sim = float(np.mean(at_sims)) if at_sims else 0.0

    if fc_pass1 > 0:
        improvement_pct = ((at_pass1 - fc_pass1) / fc_pass1) * 100
    else:
        improvement_pct = float("inf") if at_pass1 > 0 else 0.0

    mcnemar = compute_mcnemar_test(all_per_sample)

    # Token stats
    fc_tokens_list = [
        s["strategies"].get("full_context", {}).get("context_tokens", 0)
        for s in all_per_sample
    ]
    at_tokens_list = [
        s["strategies"].get("adaptive_trigger", {}).get("context_tokens", 0)
        for s in all_per_sample
        if s["strategies"].get("adaptive_trigger", {}).get("context_tokens") is not None
    ]

    results = {
        "model": model_name,
        "model_type": "open_source",
        "n_samples": total_n,
        "project": project_name,
        "project_path": project_path,
        "seed": seed,
        "timestamp": timestamp,
        "elapsed_seconds": elapsed,
        "assessment_method": "token_similarity (coverage + jaccard)",
        "strategies": {
            "full_context": {
                "passed": fc_passed,
                "total": total_n,
                "pass_at_1": fc_pass1,
                "ci_95": list(fc_ci),
                "avg_similarity": fc_avg_sim,
                "avg_context_tokens": sum(fc_tokens_list) / len(fc_tokens_list) if fc_tokens_list else 0,
            },
            "adaptive_trigger": {
                "passed": at_passed,
                "total": total_n,
                "pass_at_1": at_pass1,
                "ci_95": list(at_ci),
                "avg_similarity": at_avg_sim,
                "avg_context_tokens": sum(at_tokens_list) / len(at_tokens_list) if at_tokens_list else 0,
            },
        },
        "improvement_pct": improvement_pct,
        "mcnemar_test": mcnemar,
        "per_sample": all_per_sample,
    }

    # Save JSON
    json_path = os.path.join(results_dir, "llm_quality_opensource.json")
    with open(json_path, "w") as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print(f"  JSON saved to {json_path}")

    # Generate markdown report
    _generate_report(results, results_dir)

    # Print summary
    print(f"\n{'='*60}")
    print(f"  Model: {model_name}")
    print(f"  Samples: {total_n}")
    print(f"  Full Context:      {fc_passed}/{total_n} = {fc_pass1:.3f} pass@1  "
          f"95% CI [{fc_ci[0]:.3f}, {fc_ci[1]:.3f}]  avg_sim={fc_avg_sim:.3f}")
    print(f"  Adaptive Trigger:  {at_passed}/{total_n} = {at_pass1:.3f} pass@1  "
          f"95% CI [{at_ci[0]:.3f}, {at_ci[1]:.3f}]  avg_sim={at_avg_sim:.3f}")
    print(f"  Improvement:       {improvement_pct:+.1f}%")
    print(f"  McNemar p-value:   {mcnemar['p_value']:.4f}")
    print(f"  Time:              {elapsed:.1f}s")
    print(f"{'='*60}")

    return results


def _generate_report(results: dict, results_dir: str) -> None:
    """Generate markdown report."""
    report_path = os.path.join(results_dir, "llm_quality_opensource.md")

    fc = results["strategies"]["full_context"]
    at = results["strategies"]["adaptive_trigger"]
    mc = results["mcnemar_test"]

    lines = [
        f"# CTX LLM Downstream Quality Report (Open-Source Model)",
        "",
        f"**Date**: {results['timestamp']}",
        f"**Model**: {results['model']} (open-source, reproducible)",
        f"**Project**: {results['project']}",
        f"**Samples**: {results['n_samples']}",
        f"**Seed**: {results['seed']}",
        f"**Assessment Method**: {results['assessment_method']}",
        f"**Total Time**: {results['elapsed_seconds']:.1f}s",
        "",
        "---",
        "",
        "## Results Summary",
        "",
        "| Strategy | pass@1 | 95% CI | Passed | Total | Avg Similarity | Avg Context Tokens |",
        "|----------|--------|--------|--------|-------|----------------|-------------------|",
        f"| Full Context | {fc['pass_at_1']:.3f} | [{fc['ci_95'][0]:.3f}, {fc['ci_95'][1]:.3f}] | "
        f"{fc['passed']} | {fc['total']} | {fc['avg_similarity']:.3f} | {fc['avg_context_tokens']:.0f} |",
        f"| Adaptive Trigger | {at['pass_at_1']:.3f} | [{at['ci_95'][0]:.3f}, {at['ci_95'][1]:.3f}] | "
        f"{at['passed']} | {at['total']} | {at['avg_similarity']:.3f} | {at['avg_context_tokens']:.0f} |",
        "",
        f"**Improvement**: {results['improvement_pct']:+.1f}% (Adaptive Trigger vs Full Context)",
        "",
        "---",
        "",
        "## Statistical Significance (McNemar Test)",
        "",
        "| Metric | Value |",
        "|--------|-------|",
        f"| Chi-squared | {mc['chi2']:.4f} |",
        f"| p-value | {mc['p_value']:.4f} |",
        f"| Significant (p<0.05)? | {'Yes' if mc['p_value'] < 0.05 else 'No'} |",
        f"| Both Pass | {mc['contingency']['both_pass']} |",
        f"| FC Only Pass | {mc['contingency']['fc_only']} |",
        f"| AT Only Pass | {mc['contingency']['at_only']} |",
        f"| Both Fail | {mc['contingency']['both_fail']} |",
        "",
        "---",
        "",
        "## Methodology",
        "",
        "### Model Choice",
        f"We use **{results['model']}** -- a small, open-source code generation model",
        "that enables fully reproducible assessment without proprietary API dependencies.",
        "",
        "### Assessment Method",
        "Since small models may not generate exact matches, we use **token similarity**",
        "combining:",
        "- **Coverage**: fraction of reference code tokens present in generated code",
        "- **Jaccard similarity**: token overlap between generated and reference",
        "- **Threshold**: similarity >= 0.35 counts as pass (calibrated for small models)",
        "",
        "### Context Strategy Comparison",
        "- **Full Context**: Entire codebase concatenated (truncated to fit context window)",
        "- **Adaptive Trigger**: CTX's trigger-based selective context retrieval",
        "",
        "### Reproducibility",
        f"- Model: `{results['model']}` (HuggingFace)",
        f"- Seed: {results['seed']}",
        f"- Temperature: 0.2, top_p: 0.95",
        f"- Max new tokens: 256",
        "",
        "---",
        "",
        "## Per-Sample Results",
        "",
        "| # | Function | File | Full Context | Adaptive Trigger |",
        "|---|----------|------|-------------|-----------------|",
    ]

    for s in results["per_sample"]:
        fc_r = s["strategies"].get("full_context", {})
        at_r = s["strategies"].get("adaptive_trigger", {})

        fc_str = f"PASS ({fc_r.get('similarity', 0):.2f})" if fc_r.get("passed") else (
            f"ERROR" if fc_r.get("error") else f"FAIL ({fc_r.get('similarity', 0):.2f})"
        )
        at_str = f"PASS ({at_r.get('similarity', 0):.2f})" if at_r.get("passed") else (
            f"ERROR" if at_r.get("error") else f"FAIL ({at_r.get('similarity', 0):.2f})"
        )

        lines.append(
            f"| {s['index']+1} | `{s['function_name']}` | {s['file_path']} | {fc_str} | {at_str} |"
        )

    if fc["avg_context_tokens"] > 0:
        token_reduction = (1 - at["avg_context_tokens"] / fc["avg_context_tokens"]) * 100
        lines.extend([
            "",
            "---",
            "",
            "## Token Efficiency",
            "",
            f"- Full Context: ~{fc['avg_context_tokens']:.0f} tokens per query",
            f"- Adaptive Trigger: ~{at['avg_context_tokens']:.0f} tokens per query",
            f"- Token reduction: {token_reduction:.1f}%",
        ])

    lines.extend([
        "",
        "---",
        "",
        f"*Generated by CTX Open-Source LLM Assessment ({results['timestamp']})*",
    ])

    with open(report_path, "w") as f:
        f.write("\n".join(lines))
    print(f"  Report saved to {report_path}")


def _save_failure_report(error_msg: str, results_dir: str, timestamp: str) -> None:
    """Save a failure report when the model cannot be loaded."""
    report_path = os.path.join(results_dir, "opensource_llm_attempt.md")

    lines = [
        "# Open-Source LLM Assessment -- Attempt Report",
        "",
        f"**Date**: {timestamp}",
        f"**Status**: FAILED",
        "",
        "---",
        "",
        "## Error Details",
        "",
        f"```",
        error_msg,
        f"```",
        "",
        "## Models Attempted",
        "",
        "1. Salesforce/codegen-350M-mono (350M params) -- primary choice",
        "2. microsoft/phi-2 (2.7B params) -- fallback",
        "",
        "## Recommendations",
        "",
        "- Ensure `transformers` and `torch` are installed with CUDA support",
        "- Check GPU memory availability (RTX 3070 Ti with 8GB VRAM)",
        "",
        "---",
        "",
        f"*Generated ({timestamp})*",
    ]

    with open(report_path, "w") as f:
        f.write("\n".join(lines))
    print(f"  Failure report saved to {report_path}")


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="CTX pass@1 assessment with open-source code model"
    )
    parser.add_argument(
        "--project-path",
        default="/home/jayone/Project/GraphPrompt",
        help="Path to Python project",
    )
    parser.add_argument(
        "--n-samples", type=int, default=30,
        help="Number of function samples (default: 30)",
    )
    parser.add_argument(
        "--model",
        default="Salesforce/codegen-350M-mono",
        help="HuggingFace model name",
    )
    parser.add_argument(
        "--results-dir",
        default="benchmarks/results",
        help="Output directory",
    )
    parser.add_argument(
        "--seed", type=int, default=42,
        help="Random seed",
    )
    args = parser.parse_args()

    if not os.path.isdir(args.project_path):
        print(f"ERROR: Project path not found: {args.project_path}")
        sys.exit(1)

    run_opensource_assessment(
        project_path=args.project_path,
        n_samples=args.n_samples,
        model_name=args.model,
        results_dir=args.results_dir,
        seed=args.seed,
    )


if __name__ == "__main__":
    main()