Spaces:
Running
Running
| """Test-Time Compute pour du code : best-of-N, arbre de réparations (beam) et MCTS, avec pytest comme juge exact. | |
| Dans un MCTS classique, la valeur d'une feuille s'estime par des simulations. Ici, chaque nœud est un programme | |
| complet que pytest note exactement (part des tests réussis) : pas de simulation, l'évaluation EST le jugement. | |
| """ | |
| from __future__ import annotations | |
| import hashlib | |
| import json | |
| import math | |
| from dataclasses import dataclass, field | |
| from typing import Optional | |
| from mouh.verification.generators import parse_files | |
| from ..contracts import Candidate | |
| def signature(files: dict) -> str: | |
| return hashlib.sha256(json.dumps(files, sort_keys=True).encode()).hexdigest() | |
| def _files(text: str) -> dict: | |
| return {k: v for k, v in parse_files(text).items() if not k.startswith("test_")} | |
| class _Search: | |
| """Comptabilité commune : budget de générations, déduplication, notation.""" | |
| def __init__(self, backend, verifier, tests: dict, budget: int) -> None: | |
| self.backend, self.verifier, self.tests, self.budget = backend, verifier, tests, budget | |
| self.spent = 0 | |
| self.seen: dict[str, Candidate] = {} | |
| def done(self) -> bool: | |
| return self.spent >= self.budget or any(c.ok for c in self.seen.values()) | |
| def expand(self, messages, n: int, depth: int, origin: str, greedy_first: bool = True) -> list[Candidate]: | |
| n = max(0, min(n, self.budget - self.spent)) | |
| if n == 0: | |
| return [] | |
| out = [] | |
| for text in self.backend.complete(messages, n, greedy_first=greedy_first): | |
| self.spent += 1 | |
| files = _files(text) | |
| sig = signature(files) if files else "" | |
| if files and sig not in self.seen: | |
| c = self.verifier.score(files, self.tests) | |
| c.origin, c.depth = origin, depth | |
| self.seen[sig] = c | |
| out.append(c) | |
| return out | |
| def ranked(self) -> list[Candidate]: | |
| return sorted(self.seen.values(), key=Candidate.key, reverse=True) | |
| class BestOfN: | |
| """Une génération gloutonne + des échantillons ; la référence à battre, à budget égal.""" | |
| name = "best_of_n" | |
| def solve(self, request, plan, backend, verifier, build) -> list[Candidate]: | |
| tests = dict(request.task.test_files) | |
| if not tests: # rien pour juger : une seule proposition, comme la v3.8 | |
| files = _files(backend.complete(build(), 1)[0]) | |
| return [Candidate(files, "racine")] if files else [] | |
| s = _Search(backend, verifier, tests, plan.budget) | |
| s.expand(build(), plan.budget, 0, "racine") | |
| return s.ranked() | |
| class RepairBeam: | |
| """Arbre de réparations en faisceau : à chaque niveau, les `width` meilleurs sont réparés avec leur trace.""" | |
| name = "beam" | |
| def __init__(self, width: int = 2, depth: int = 3) -> None: | |
| self.width, self.depth = width, depth | |
| def solve(self, request, plan, backend, verifier, build) -> list[Candidate]: | |
| tests = dict(request.task.test_files) | |
| if not tests: | |
| return BestOfN().solve(request, plan, backend, verifier, build) | |
| s = _Search(backend, verifier, tests, plan.budget) | |
| frontier = s.expand(build(), max(1, plan.budget // 2), 0, "racine") | |
| for d in range(1, self.depth + 1): | |
| if s.done or not frontier: | |
| break | |
| parents = sorted(frontier, key=Candidate.key, reverse=True)[: self.width] | |
| frontier = [] | |
| for p in parents: | |
| frontier += s.expand(build(previous=p.files, diagnostics=[p.output]), 1, d, "reparation") | |
| return s.ranked() | |
| class _Node: | |
| cand: Optional[Candidate] # None pour la racine (l'énoncé seul) | |
| parent: Optional["_Node"] = None | |
| children: list = field(default_factory=list) | |
| visits: int = 0 | |
| total: float = 0.0 | |
| expansions: int = 0 # générations déjà tirées depuis ce nœud | |
| def q(self) -> float: | |
| return self.total / self.visits if self.visits else 0.0 | |
| class RepairMCTS: | |
| """MCTS (sélection UCT) sur l'arbre des réparations. | |
| - Sélection : parmi la racine (nouvelle proposition depuis l'énoncé) et tous les candidats déjà notés, | |
| on choisit celui qui maximise Q + c * sqrt(ln N / (1 + n)). | |
| - Expansion : la racine produit un nouvel échantillon ; un candidat est réparé avec sa propre sortie pytest | |
| (1re réparation gloutonne, les suivantes échantillonnées). | |
| - Évaluation : pytest, exact (part des tests réussis). | |
| - Rétropropagation : la valeur remonte jusqu'à la racine. | |
| Arrêt au premier candidat qui passe tous les tests ou quand le budget est dépensé. | |
| """ | |
| name = "mcts" | |
| def __init__(self, c: float = 1.0, root_samples: int = 2, max_children: int = 3) -> None: | |
| self.c, self.root_samples, self.max_children = c, root_samples, max_children | |
| def solve(self, request, plan, backend, verifier, build) -> list[Candidate]: | |
| tests = dict(request.task.test_files) | |
| if not tests: | |
| return BestOfN().solve(request, plan, backend, verifier, build) | |
| s = _Search(backend, verifier, tests, plan.budget) | |
| root = _Node(None) | |
| nodes: list[_Node] = [] | |
| def attach(parent: _Node, cands: list[Candidate]) -> None: | |
| for c in cands: | |
| node = _Node(c, parent) | |
| parent.children.append(node) | |
| nodes.append(node) | |
| v, cur = c.value, node | |
| while cur is not None: | |
| cur.visits += 1 | |
| cur.total += v | |
| cur = cur.parent | |
| first = s.expand(build(), min(self.root_samples, plan.budget), 0, "racine") | |
| root.expansions += min(self.root_samples, plan.budget) | |
| attach(root, first) | |
| while not s.done: | |
| pool = [root] + [n for n in nodes if n.expansions < self.max_children and not n.cand.ok] | |
| total = max(1, root.visits) | |
| def uct(n: _Node) -> float: # Q = moyenne du sous-arbre (le nœud compris) | |
| return n.q + self.c * math.sqrt(math.log(total + 1) / (1 + n.expansions)) | |
| target = max(pool, key=uct) | |
| if target is root: | |
| found = s.expand(build(), 1, 0, "racine", greedy_first=False) | |
| else: | |
| found = s.expand(build(previous=target.cand.files, diagnostics=[target.cand.output]), 1, | |
| target.cand.depth + 1, "reparation", greedy_first=target.expansions == 0) | |
| target.expansions += 1 | |
| if found: | |
| attach(target, found) | |
| else: # doublon ou sortie illisible : la visite compte quand même, pour ne pas reboucler | |
| cur = target | |
| while cur is not None: | |
| cur.visits += 1 | |
| cur = cur.parent | |
| return s.ranked() | |
| class AdaptiveStrategy: | |
| """Choisit la stratégie selon le plan du routeur : le raisonnement hybride.""" | |
| name = "adaptatif" | |
| def __init__(self) -> None: | |
| self.by_mode = {"direct": BestOfN(), "arbre": RepairBeam(), "mcts": RepairMCTS()} | |
| def solve(self, request, plan, backend, verifier, build) -> list[Candidate]: | |
| return self.by_mode.get(plan.mode, self.by_mode["direct"]).solve(request, plan, backend, verifier, build) | |