colgora
141df5f04e
- Nouvelle commande `inkflow benchmark` : compare la sortie d'analyse aux fichiers de référence (data/<slug>/reference/), met plusieurs modèles en concurrence, table rich + rapport JSON. Métriques : attribution de locuteur, incises, type/glued. Flags --models, --temperature, --reasoning, --stream, --use-cached + suivi par chapitre. - analysis/benchmark.py : scoring pur (testable) + runner multi-modèles (un MLX à la fois). - gemma.py : support des modèles à raisonnement (retrait de la pensée, désactivation via enable_thinking hors --reasoning, arrêt anticipé sur JSON complet, plafond + température dédiés anti-boucle), récupération du chat_template manquant (fix Mistral), streaming des tokens (set_token_sink). - settings.py : gemma_reasoning, gemma_reasoning_max_tokens, gemma_reasoning_temperature. - Tests : test_benchmark.py (scoring pur), test_gemma_reasoning.py. Conclusion benchmark : Qwen3.6-27B-8bit non-raisonnant = meilleur modèle d'analyse.
252 lines
10 KiB
Python
252 lines
10 KiB
Python
"""Wrapper mlx-lm autour de Gemma pour l'analyse de texte.
|
|
|
|
Charge le modele paresseusement (une seule fois par process) et expose des
|
|
helpers de generation, dont un `generate_json` tolerant qui extrait le premier
|
|
objet/array JSON valide de la sortie du modele.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from functools import lru_cache
|
|
from typing import Any, Optional
|
|
|
|
from ..settings import get_settings
|
|
|
|
# Bornes d'un bloc JSON dans une reponse potentiellement bavarde.
|
|
_JSON_SPAN_RE = re.compile(r"(\{.*\}|\[.*\])", re.DOTALL)
|
|
_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL)
|
|
|
|
# Marqueurs de FIN de chaine de pensee : on ne garde que ce qui suit le dernier.
|
|
# - balises type DeepSeek-R1 / Qwen-think
|
|
# - format a canaux type Gemma 4 / Harmony (la pensee est close par <channel|>)
|
|
_REASONING_END_MARKERS = ("</think>", "<channel|>", "<|channel|>")
|
|
# Prefixe de canal/think non ferme reste en tete (pensee tronquee) : a retirer.
|
|
_REASONING_OPEN_RE = re.compile(r"^\s*(?:<\|?channel\|?>\s*\w*|<think>)", re.IGNORECASE)
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def _load(model_id: str):
|
|
# Import paresseux : evite de charger mlx tant qu'on n'analyse pas.
|
|
from mlx_lm import load
|
|
return load(model_id)
|
|
|
|
|
|
# Hook de streaming optionnel. Si defini, `generate()` diffuse chaque morceau de
|
|
# texte AU FIL de la generation (pensee comprise, avant tout nettoyage) en
|
|
# appelant ce callback. Utilise par `inkflow benchmark --stream` pour voir les
|
|
# tokens en temps reel. None -> generation par lot classique (plus rapide).
|
|
_TOKEN_SINK: Optional[Any] = None
|
|
|
|
|
|
def set_token_sink(callback) -> None:
|
|
"""Definit (ou retire avec None) le callback de streaming des tokens."""
|
|
global _TOKEN_SINK
|
|
_TOKEN_SINK = callback
|
|
|
|
|
|
def _resolve_chat_template(model_id: str, tokenizer) -> Optional[str]:
|
|
"""Renvoie un template de chat a passer explicitement, ou None.
|
|
|
|
Certaines conversions (Mistral recents...) logent leur template dans un
|
|
fichier `chat_template.jinja` que le downloader de mlx-lm n'embarque pas
|
|
toujours : `tokenizer.chat_template` est alors vide et `apply_chat_template`
|
|
echoue. On recupere alors le fichier officiel du repo. None si le tokenizer
|
|
possede deja un template (cas courant) ou si aucun n'est disponible.
|
|
"""
|
|
if getattr(tokenizer, "chat_template", None):
|
|
return None
|
|
from pathlib import Path
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
# Selon les conversions : fichier Jinja brut, ou JSON {"chat_template": ...}.
|
|
for fname in ("chat_template.jinja", "chat_template.json"):
|
|
try:
|
|
text = Path(hf_hub_download(model_id, fname)).read_text(encoding="utf-8")
|
|
except Exception: # noqa: BLE001 — fichier absent, on tente le suivant
|
|
continue
|
|
if fname.endswith(".json"):
|
|
data = json.loads(text)
|
|
return data.get("chat_template") if isinstance(data, dict) else None
|
|
return text
|
|
return None # aucun template dispo -> apply_chat_template levera une erreur claire
|
|
|
|
|
|
class Gemma:
|
|
"""Petite facade autour de mlx-lm pour piloter Gemma."""
|
|
|
|
def __init__(self, model_id: Optional[str] = None):
|
|
self.model_id = model_id or get_settings().gemma_model
|
|
self._model = None
|
|
self._tokenizer = None
|
|
self._chat_template = None # template recupere si absent du tokenizer
|
|
|
|
def _ensure_loaded(self) -> None:
|
|
if self._model is None:
|
|
self._model, self._tokenizer = _load(self.model_id)
|
|
self._chat_template = _resolve_chat_template(
|
|
self.model_id, self._tokenizer)
|
|
|
|
def generate(
|
|
self,
|
|
prompt: str,
|
|
*,
|
|
system: Optional[str] = None,
|
|
max_tokens: Optional[int] = None,
|
|
temperature: Optional[float] = None,
|
|
) -> str:
|
|
"""Genere une reponse texte a partir d'un prompt (template de chat).
|
|
|
|
`max_tokens`/`temperature` non fournis -> valeurs des reglages courants.
|
|
"""
|
|
self._ensure_loaded()
|
|
settings = get_settings()
|
|
if max_tokens is None:
|
|
max_tokens = settings.gemma_max_tokens
|
|
# En mode raisonnement, plafond dedie (garde-fou anti-boucle) ; la
|
|
# generation s'arrete de toute facon des que le JSON post-pensee est
|
|
# complet (cf. boucle de streaming ci-dessous).
|
|
if settings.gemma_reasoning:
|
|
max_tokens = max(max_tokens, settings.gemma_reasoning_max_tokens)
|
|
if temperature is None:
|
|
temperature = settings.gemma_temperature
|
|
# Decodage glouton (temp 0) + raisonnement = boucles de pensee sans fin.
|
|
# On force un echantillonnage minimal en mode raisonnement.
|
|
if settings.gemma_reasoning and temperature == 0.0:
|
|
temperature = settings.gemma_reasoning_temperature
|
|
from mlx_lm.sample_utils import make_sampler
|
|
|
|
messages = []
|
|
if system:
|
|
messages.append({"role": "system", "content": system})
|
|
messages.append({"role": "user", "content": prompt})
|
|
# Modeles hybrides (Qwen3...) : hors mode raisonnement, on DESACTIVE la
|
|
# pensee via enable_thinking=False -> JSON direct, bien plus rapide. Avec
|
|
# --reasoning, on laisse penser puis on retire la pensee en aval. Ce
|
|
# kwarg est ignore par les templates qui ne l'utilisent pas (Gemma...).
|
|
template_kwargs = {}
|
|
if not settings.gemma_reasoning:
|
|
template_kwargs["enable_thinking"] = False
|
|
formatted = self._tokenizer.apply_chat_template(
|
|
messages, add_generation_prompt=True, tokenize=False,
|
|
chat_template=self._chat_template, # None -> celui du tokenizer
|
|
**template_kwargs,
|
|
)
|
|
sampler = make_sampler(temp=temperature)
|
|
# On streame (token par token) si : un sink est branche (--stream) OU on
|
|
# est en mode raisonnement (pour pouvoir s'arreter des que la reponse est
|
|
# prete, sans subir les boucles de pensee sans fin). Sinon, lot rapide.
|
|
if _TOKEN_SINK is not None or settings.gemma_reasoning:
|
|
from mlx_lm import stream_generate
|
|
parts = []
|
|
seen_end = False # marqueur de fin de pensee rencontre
|
|
for resp in stream_generate(
|
|
self._model, self._tokenizer, prompt=formatted,
|
|
max_tokens=max_tokens, sampler=sampler,
|
|
):
|
|
parts.append(resp.text)
|
|
if _TOKEN_SINK is not None:
|
|
_TOKEN_SINK(resp.text)
|
|
# Arret anticipe : une fois la pensee close, des que le JSON
|
|
# post-pensee est complet, inutile de continuer a generer.
|
|
if settings.gemma_reasoning and ("}" in resp.text or "]" in resp.text):
|
|
buf = "".join(parts)
|
|
if not seen_end:
|
|
seen_end = any(mk in buf for mk in _REASONING_END_MARKERS)
|
|
if seen_end and _has_complete_json(_strip_reasoning(buf)):
|
|
break
|
|
if _TOKEN_SINK is not None:
|
|
_TOKEN_SINK("\n") # separe les generations successives
|
|
raw = "".join(parts)
|
|
else:
|
|
from mlx_lm import generate
|
|
raw = generate(
|
|
self._model,
|
|
self._tokenizer,
|
|
prompt=formatted,
|
|
max_tokens=max_tokens,
|
|
sampler=sampler,
|
|
verbose=False,
|
|
)
|
|
# Retire la chaine de pensee des modeles a raisonnement (sinon des
|
|
# fragments de la "pensee" parasitent l'extraction JSON en aval).
|
|
if settings.gemma_reasoning:
|
|
return _strip_reasoning(raw)
|
|
return raw
|
|
|
|
def generate_json(
|
|
self,
|
|
prompt: str,
|
|
*,
|
|
system: Optional[str] = None,
|
|
max_tokens: Optional[int] = None,
|
|
temperature: Optional[float] = None,
|
|
retries: int = 1,
|
|
) -> Any:
|
|
"""Genere puis parse un JSON. Reessaie en cas d'echec de parsing.
|
|
|
|
`max_tokens`/`temperature` non fournis -> valeurs des reglages courants.
|
|
"""
|
|
last_err: Optional[Exception] = None
|
|
for attempt in range(retries + 1):
|
|
raw = self.generate(
|
|
prompt, system=system, max_tokens=max_tokens,
|
|
temperature=temperature if attempt == 0 else 0.0,
|
|
)
|
|
try:
|
|
return _extract_json(raw)
|
|
except Exception as exc: # noqa: BLE001
|
|
last_err = exc
|
|
raise ValueError(f"Reponse JSON invalide apres {retries + 1} essais: {last_err}")
|
|
|
|
|
|
def _strip_reasoning(text: str) -> str:
|
|
"""Retire la chaine de pensee d'un modele a raisonnement.
|
|
|
|
Ne conserve que ce qui suit le dernier marqueur de fin de pensee
|
|
(`</think>`, `<channel|>`...). Si seul un marqueur d'ouverture non ferme
|
|
subsiste (pensee tronquee par le budget de tokens), on le retire en tete
|
|
pour eviter de parser la pensee a la place de la reponse.
|
|
"""
|
|
t = text
|
|
for marker in _REASONING_END_MARKERS:
|
|
if marker in t:
|
|
t = t.rsplit(marker, 1)[-1]
|
|
t = _REASONING_OPEN_RE.sub("", t)
|
|
return t.strip()
|
|
|
|
|
|
def _has_complete_json(text: str) -> bool:
|
|
"""True si `text` contient deja un objet/array JSON complet et parsable.
|
|
|
|
Sert a stopper la generation des modeles a raisonnement des que la reponse
|
|
finale est ecrite (evite de consommer le budget en boucles de pensee).
|
|
"""
|
|
try:
|
|
_extract_json(text)
|
|
return True
|
|
except Exception: # noqa: BLE001
|
|
return False
|
|
|
|
|
|
def _extract_json(text: str) -> Any:
|
|
"""Extrait le premier objet/array JSON d'une reponse libre du modele.
|
|
|
|
Tolere le texte parasite avant/apres (y compris un 2e bloc) grace a
|
|
raw_decode, qui s'arrete au premier JSON complet.
|
|
"""
|
|
text = text.strip()
|
|
fence = _FENCE_RE.search(text)
|
|
if fence:
|
|
text = fence.group(1).strip()
|
|
decoder = json.JSONDecoder()
|
|
# Cherche le 1er debut de structure JSON et decode a partir de la.
|
|
for i, ch in enumerate(text):
|
|
if ch in "[{":
|
|
try:
|
|
obj, _ = decoder.raw_decode(text[i:])
|
|
return obj
|
|
except json.JSONDecodeError:
|
|
continue
|
|
raise ValueError("aucun JSON trouve dans la reponse")
|