evaluate_model.py

#!/usr/bin/env python3
"""
training/scripts/evaluate_model.py
Évalue le modèle fine-tuné sur les 6 tâches VORTEX GOD.

Métriques par tâche :
  planner   → json_valid_rate, steps_count_correct, risks_present
  engineer  → syntax_valid_rate, exec_pass_rate, has_docstring
  critic    → json_valid_rate, verdict_present, issues_severity_present
  researcher→ has_sections, has_citations, has_lacunes
  scientist → json_valid_rate, hypothesis_length, has_experiment
  optimizer → json_valid_rate, has_changes, speedup_numeric

Score global = moyenne pondérée (0.0 – 1.0)

Usage :
  python training/scripts/evaluate_model.py \
    --model training/checkpoints/vortex-lora \
    --dataset training/data/vortex_dataset.jsonl \
    --base-model Qwen/Qwen3-8B-Instruct \
    --n 50 \
    --output training/eval_results.json
"""

from __future__ import annotations

import argparse
import ast
import json
import logging
import os
import random
import re
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

log = logging.getLogger("evaluate_model")
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


# ─────────────────────────────────────────────
# Métriques par tâche
# ─────────────────────────────────────────────

def _is_valid_json(text: str) -> bool:
    clean = re.sub(r"```(?:json)?\n?", "", text).replace("```", "").strip()
    for pattern in (r"\{[\s\S]*\}", r"\[[\s\S]*\]"):
        m = re.search(pattern, clean)
        if m:
            try:
                json.loads(m.group(0))
                return True
            except Exception:
                pass
    return False


def _is_valid_python(text: str) -> bool:
    clean = re.sub(r"```(?:python|py)?\n?", "", text).replace("```", "").strip()
    try:
        ast.parse(clean)
        return len(clean) > 30
    except SyntaxError:
        return False


def _exec_python(text: str, timeout: float = 5.0) -> bool:
    """Exécute le code dans un sous-process isolé. Retourne True si exit code 0."""
    import subprocess, tempfile
    clean = re.sub(r"```(?:python|py)?\n?", "", text).replace("```", "").strip()
    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
        f.write(clean)
        fname = f.name
    try:
        result = subprocess.run(
            ["python3", fname],
            timeout=timeout,
            capture_output=True,
        )
        return result.returncode == 0
    except Exception:
        return False
    finally:
        os.unlink(fname)


TASK_METRICS: Dict[str, List[Tuple[str, Callable[[str], float], float]]] = {
    # (metric_name, scorer_fn returns 0.0-1.0, weight)
    "planner": [
        ("json_valid",      lambda r: 1.0 if _is_valid_json(r) else 0.0,                                  2.0),
        ("has_steps",       lambda r: 1.0 if '"steps"' in r and '"id"' in r else 0.0,                    1.5),
        ("has_goal",        lambda r: 1.0 if '"goal"' in r else 0.0,                                      1.0),
        ("has_risks",       lambda r: 1.0 if '"risks"' in r else 0.0,                                     0.5),
        ("steps_count",     lambda r: min(1.0, _count_steps(r) / 4),                                      1.0),
    ],
    "engineer": [
        ("syntax_valid",    lambda r: 1.0 if _is_valid_python(r) else 0.0,                                3.0),
        ("has_imports",     lambda r: 1.0 if re.search(r"^import |^from ", r, re.MULTILINE) else 0.0,    1.0),
        ("has_docstring",   lambda r: 1.0 if '"""' in r or "'''" in r else 0.0,                          0.5),
        ("has_type_hints",  lambda r: 1.0 if " -> " in r or ": " in r else 0.0,                         0.5),
        ("exec_pass",       lambda r: 1.0 if _exec_python(r) else 0.0,                                   2.0),
    ],
    "critic": [
        ("json_valid",      lambda r: 1.0 if _is_valid_json(r) else 0.0,                                  2.0),
        ("has_verdict",     lambda r: 1.0 if '"verdict"' in r else 0.0,                                   1.5),
        ("has_issues",      lambda r: 1.0 if '"issues"' in r else 0.0,                                    1.5),
        ("has_severity",    lambda r: 1.0 if '"severity"' in r else 0.0,                                  1.0),
        ("has_fix",         lambda r: 1.0 if '"fix"' in r else 0.0,                                      0.5),
    ],
    "researcher": [
        ("has_sections",    lambda r: min(1.0, r.count("##") / 2),                                        2.0),
        ("has_citations",   lambda r: 1.0 if "[source:" in r.lower() else 0.0,                            2.0),
        ("has_lacunes",     lambda r: 1.0 if "lacune" in r.lower() or "manque" in r.lower() else 0.0,    1.0),
        ("length_ok",       lambda r: 1.0 if 200 <= len(r.split()) <= 600 else 0.5,                       1.0),
    ],
    "scientist": [
        ("json_valid",      lambda r: 1.0 if _is_valid_json(r) else 0.0,                                  2.0),
        ("has_hypothesis",  lambda r: 1.0 if '"hypothesis"' in r else 0.0,                                2.0),
        ("has_experiment",  lambda r: 1.0 if '"experiment"' in r else 0.0,                                1.5),
        ("has_confidence",  lambda r: 1.0 if '"confidence"' in r else 0.0,                                0.5),
        ("has_null_hyp",    lambda r: 1.0 if '"null_hypothesis"' in r else 0.0,                          0.5),
    ],
    "optimizer": [
        ("json_valid",      lambda r: 1.0 if _is_valid_json(r) else 0.0,                                  2.0),
        ("has_optimized",   lambda r: 1.0 if '"optimized"' in r else 0.0,                                 2.0),
        ("has_changes",     lambda r: 1.0 if '"changes"' in r else 0.0,                                   1.5),
        ("has_speedup",     lambda r: 1.0 if '"expected_speedup"' in r else 0.0,                          0.5),
        ("has_rationale",   lambda r: 1.0 if '"rationale"' in r else 0.0,                                 0.5),
    ],
}

TASK_WEIGHTS = {
    "planner":    1.0,
    "engineer":   1.5,   # tâche la plus critique
    "critic":     1.0,
    "researcher": 0.8,
    "scientist":  1.0,
    "optimizer":  0.8,
}


def _count_steps(text: str) -> int:
    clean = re.sub(r"```(?:json)?\n?", "", text).replace("```", "").strip()
    m = re.search(r"\{[\s\S]*\}", clean)
    if m:
        try:
            data = json.loads(m.group(0))
            return len(data.get("steps", []))
        except Exception:
            pass
    return 0


def score_response(task_type: str, response: str) -> Dict[str, float]:
    """Calcule les métriques pour une réponse."""
    metrics_def = TASK_METRICS.get(task_type, [])
    scores = {}
    total_weight = 0.0
    weighted_sum = 0.0

    for metric_name, scorer, weight in metrics_def:
        try:
            s = scorer(response)
        except Exception:
            s = 0.0
        scores[metric_name] = round(s, 3)
        weighted_sum  += s * weight
        total_weight  += weight

    scores["task_score"] = round(weighted_sum / total_weight, 4) if total_weight > 0 else 0.0
    return scores


# ─────────────────────────────────────────────
# Génération avec le modèle évalué
# ─────────────────────────────────────────────

class ModelEvaluator:
    def __init__(self, model_path: str, base_model: str, use_gguf: bool = False):
        self.model_path = model_path
        self.base_model = base_model
        self.use_gguf   = use_gguf or model_path.endswith(".gguf")
        self._model     = None
        self._tokenizer = None
        self._llama     = None
        self._load()

    def _load(self):
        if self.use_gguf:
            from llama_cpp import Llama
            self._llama = Llama(model_path=self.model_path, n_ctx=2048, n_gpu_layers=35, verbose=False)
            log.info(f"[Eval] GGUF chargé : {self.model_path}")
        else:
            import torch
            from peft import PeftModel
            from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

            log.info(f"[Eval] Chargement LoRA depuis : {self.model_path}")
            bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                     bnb_4bit_compute_dtype=torch.bfloat16)
            self._tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True)
            base = AutoModelForCausalLM.from_pretrained(
                self.base_model,
                quantization_config = bnb,
                device_map          = "auto",
                trust_remote_code   = True,
                torch_dtype         = torch.bfloat16,
            )
            self._model = PeftModel.from_pretrained(base, self.model_path)
            self._model.eval()
            log.info(f"[Eval] Modèle LoRA chargé ({self._model.num_parameters():,} params)")

    def generate(self, system: str, user: str, max_tokens: int = 512, temperature: float = 0.1) -> str:
        if self.use_gguf and self._llama:
            out = self._llama.create_chat_completion(
                messages=[{"role": "system", "content": system},
                          {"role": "user", "content": user}],
                max_tokens=max_tokens, temperature=max(temperature, 0.01),
            )
            return out["choices"][0]["message"]["content"].strip()

        import torch
        messages = [{"role": "system", "content": system}, {"role": "user", "content": user}]
        text     = self._tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs   = self._tokenizer(text, return_tensors="pt").to(self._model.device)
        with torch.no_grad():
            out = self._model.generate(
                **inputs, max_new_tokens=max_tokens, temperature=temperature,
                do_sample=True, pad_token_id=self._tokenizer.eos_token_id,
            )
        new_ids = out[0][inputs.input_ids.shape[1]:]
        return self._tokenizer.decode(new_ids, skip_special_tokens=True).strip()


# ─────────────────────────────────────────────
# Pipeline d'évaluation
# ─────────────────────────────────────────────

def run_evaluation(args: argparse.Namespace) -> Dict[str, Any]:
    # Charger les exemples de test depuis le dataset
    from build_dataset import SYSTEMS

    all_examples: List[Dict] = []
    with open(args.dataset) as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    all_examples.append(json.loads(line))
                except Exception:
                    pass

    # Sélectionner n exemples équilibrés par tâche
    by_task: Dict[str, List] = {}
    for ex in all_examples:
        t = ex.get("task_type", "unknown")
        by_task.setdefault(t, []).append(ex)

    test_examples = []
    n_per_task = max(1, args.n // len(SYSTEMS))
    for task, exs in by_task.items():
        test_examples.extend(random.sample(exs, min(n_per_task, len(exs))))

    log.info(f"Évaluation sur {len(test_examples)} exemples")

    # Charger le modèle
    evaluator = ModelEvaluator(args.model, args.base_model, use_gguf=args.gguf)

    # Évaluation
    results_by_task: Dict[str, List[Dict]] = {}
    total_latency = 0.0

    for i, ex in enumerate(test_examples):
        task_type = ex.get("task_type", "unknown")
        messages  = ex["messages"]
        system    = messages[0]["content"]
        user      = messages[1]["content"]
        reference = messages[2]["content"]

        t0       = time.time()
        response = evaluator.generate(system, user, max_tokens=600, temperature=0.1)
        latency  = time.time() - t0
        total_latency += latency

        scores = score_response(task_type, response)
        scores["latency_s"] = round(latency, 2)

        # Comparaison avec la référence (longueur relative)
        ref_len  = len(reference.split())
        resp_len = len(response.split())
        scores["length_ratio"] = round(resp_len / max(ref_len, 1), 2)

        results_by_task.setdefault(task_type, []).append({
            "user":     user[:150],
            "response": response[:400],
            "scores":   scores,
        })

        if (i + 1) % 10 == 0:
            log.info(f"  Progression : {i+1}/{len(test_examples)}")

    # Agrégation
    summary: Dict[str, Any] = {
        "model":         args.model,
        "base_model":    args.base_model,
        "n_examples":    len(test_examples),
        "avg_latency_s": round(total_latency / len(test_examples), 2),
        "tasks":         {},
        "global_score":  0.0,
    }

    weighted_scores = []
    for task_type, results in results_by_task.items():
        task_scores = [r["scores"]["task_score"] for r in results]
        avg_score   = sum(task_scores) / len(task_scores)
        weight      = TASK_WEIGHTS.get(task_type, 1.0)

        # Détail par métrique
        metric_avgs = {}
        for metric_name, _, _ in TASK_METRICS.get(task_type, []):
            vals = [r["scores"].get(metric_name, 0.0) for r in results]
            metric_avgs[metric_name] = round(sum(vals) / len(vals), 3)

        summary["tasks"][task_type] = {
            "n":           len(results),
            "avg_score":   round(avg_score, 4),
            "weight":      weight,
            "metrics":     metric_avgs,
        }
        weighted_scores.append(avg_score * weight)
        log.info(f"  [{task_type:12s}] score={avg_score:.4f}  n={len(results)}")

    total_weight = sum(TASK_WEIGHTS.get(t, 1.0) for t in results_by_task)
    summary["global_score"] = round(sum(weighted_scores) / total_weight, 4) if total_weight else 0.0

    log.info(f"\n{'='*50}")
    log.info(f"Score global : {summary['global_score']:.4f}")
    log.info(f"Latence moy. : {summary['avg_latency_s']:.2f}s")
    log.info(f"{'='*50}")

    # Sauvegarde
    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2))
    log.info(f"✅ Résultats → {args.output}")

    return summary


def main():
    parser = argparse.ArgumentParser(description="Évaluation du modèle VORTEX GOD")
    parser.add_argument("--model",      required=True, help="Chemin checkpoint LoRA ou .gguf")
    parser.add_argument("--dataset",    default="training/data/full_dataset.jsonl")
    parser.add_argument("--base-model", default="Qwen/Qwen3-8B-Instruct")
    parser.add_argument("--n",          type=int, default=50, help="Nb exemples de test")
    parser.add_argument("--output",     default="training/eval_results.json")
    parser.add_argument("--gguf",       action="store_true", help="Utiliser le mode GGUF")
    args = parser.parse_args()

    run_evaluation(args)


if __name__ == "__main__":
    main()