t2a-finetune/scripts/12_generate_pipeline_examples.py

#!/usr/bin/env python3
"""
Génère des exemples d'entraînement au format RÉEL du pipeline T2A
à partir du cache Ollama (gemma3:12b = gold standard).

Le cache contient les paires (diagnostic → code + raisonnement)
produites par gemma3:12b sur les 250 dossiers. On reconstruit
des prompts proches du pipeline et on utilise les réponses du cache
comme labels supervisés.

V2 : Utilise le cache actuel (1 840 entrées vs 100).
     3 templates : court + long + CPAM contre-argumentation.
     Tous les textes génèrent une version courte ET longue.
     Cible ~4 000 exemples (2×1840 + CPAM bonus).

Produit : data/processed/pipeline_chatml.jsonl

Usage :
  python scripts/12_generate_pipeline_examples.py
"""

import json
import random
from pathlib import Path

random.seed(42)

BASE = Path(__file__).resolve().parent.parent
T2A = Path("/home/dom/ai/t2a")
CACHE_PATH = T2A / "data" / "ollama_cache.json"
CACHE_BACKUP = T2A / "data" / "ollama_cache_gemma3.bak"
CIM10_DICT = T2A / "data" / "cim10_dict.json"
CIM10_SUPP = T2A / "data" / "cim10_supplements.json"
OUTPUT = BASE / "data" / "processed" / "pipeline_chatml.jsonl"


def load_cim10_dict() -> dict[str, str]:
    """Charge le dictionnaire CIM-10 (code → libellé)."""
    d = {}
    if CIM10_DICT.exists():
        d.update(json.loads(CIM10_DICT.read_text()))
    if CIM10_SUPP.exists():
        d.update(json.loads(CIM10_SUPP.read_text()))
    return d


def load_cache_entries() -> dict:
    """Charge toutes les entrées du cache (actuel + backup)."""
    entries = {}
    for path in [CACHE_PATH, CACHE_BACKUP]:
        if path.exists():
            data = json.loads(path.read_text())
            new = sum(1 for k in data.get("entries", {}) if k not in entries)
            entries.update(data.get("entries", {}))
            print(f"  {path.name}: {len(data.get('entries', {}))} entrées (+{new} nouvelles)")
    return entries


SYSTEM_PROMPT = (
    "Tu es un médecin DIM (Département d'Information Médicale) expert en codage PMSI. "
    "Tu codes les diagnostics en CIM-10 en suivant une démarche structurée : "
    "analyse clinique, identification des codes candidats, discrimination, "
    "vérification des règles PMSI."
)

# Template simplifié reproduisant la structure du prompt pipeline
# (sans les sources RAG qui ne sont pas dans le cache)
PROMPT_TEMPLATE_DP = """Code ce diagnostic en CIM-10 pour le PMSI.

RÈGLES IMPÉRATIVES :
- Privilégie le code le plus SPÉCIFIQUE disponible (4e ou 5e caractère)
- Vérifie les notes d'inclusion/exclusion de chaque code candidat
- Le DP doit refléter le motif principal de prise en charge du séjour
- EXCLUSION SYMPTÔME : Si le diagnostic est un symptôme (R00-R99) et qu'un diagnostic précis existe, le symptôme ne doit PAS être codé comme DP

DIAGNOSTIC À CODER : "{texte}"
TYPE : DP (diagnostic principal)

Réponds UNIQUEMENT avec un objet JSON :
{{
  "analyse_clinique": "que signifie ce diagnostic sur le plan médical",
  "codes_candidats": "quels codes CIM-10 sont compatibles",
  "discrimination": "pourquoi choisir ce code plutôt qu'un autre",
  "regle_pmsi": "conformité aux règles PMSI pour un DP",
  "code": "X99.9",
  "confidence": "high ou medium ou low",
  "justification": "explication courte en français"
}}"""

PROMPT_TEMPLATE_DAS = """Code ce diagnostic en CIM-10 pour le PMSI.

RÈGLES IMPÉRATIVES :
- Privilégie le code le plus SPÉCIFIQUE disponible (4e ou 5e caractère)
- Vérifie les notes d'inclusion/exclusion de chaque code candidat
- Un DAS doit avoir mobilisé des ressources supplémentaires pendant le séjour
- EXCLUSION SYMPTÔME : Si le diagnostic est un symptôme (R00-R99) et qu'un diagnostic précis existe, le symptôme ne doit PAS être codé comme DAS

DIAGNOSTIC À CODER : "{texte}"
TYPE : DAS (diagnostic associé significatif)

Réponds UNIQUEMENT avec un objet JSON :
{{
  "analyse_clinique": "que signifie ce diagnostic sur le plan médical",
  "codes_candidats": "quels codes CIM-10 sont compatibles",
  "discrimination": "pourquoi choisir ce code plutôt qu'un autre",
  "regle_pmsi": "conformité aux règles PMSI pour un DAS",
  "code": "X99.9",
  "confidence": "high ou medium ou low",
  "justification": "explication courte en français"
}}"""

# Version longue avec contexte patient (simulé) pour entraîner sur des prompts longs
PROMPT_TEMPLATE_DP_LONG = """Code ce diagnostic en CIM-10 pour le PMSI.

RÈGLES IMPÉRATIVES :
- Le code doit provenir UNIQUEMENT de la nomenclature CIM-10 FR 2026
- Distingue la DESCRIPTION CLINIQUE (ce que le médecin écrit) de la LOGIQUE DE CODAGE (ce que l'ATIH impose)
- Privilégie le code le plus SPÉCIFIQUE disponible (4e ou 5e caractère)
- Vérifie les notes d'inclusion/exclusion de chaque code candidat
- Si le diagnostic est un DP, il doit refléter le motif principal de prise en charge du séjour
- EXCLUSION SYMPTÔME : Si le diagnostic est un symptôme (R00-R99) et qu'un diagnostic précis (Chapitres I-XIV, A00-N99) expliquant ce symptôme est présent, le symptôme ne doit PAS être codé

DIAGNOSTIC À CODER : "{texte}"
TYPE : DP (diagnostic principal)

CONTEXTE CLINIQUE :
{contexte}

SOURCES CIM-10 :
{sources}

Réponds UNIQUEMENT avec un objet JSON au format suivant, sans aucun texte avant ou après :
{{
  "analyse_clinique": "que signifie ce diagnostic sur le plan médical",
  "codes_candidats": "quels codes CIM-10 des sources sont compatibles",
  "discrimination": "pourquoi choisir ce code plutôt qu'un autre (inclusions/exclusions, spécificité)",
  "regle_pmsi": "conformité aux règles PMSI pour un DP (guide méthodologique)",
  "code": "X99.9",
  "confidence": "high ou medium ou low",
  "justification": "explication courte en français"
}}"""

PROMPT_TEMPLATE_DAS_LONG = """Code ce diagnostic en CIM-10 pour le PMSI.

RÈGLES IMPÉRATIVES :
- Le code doit provenir UNIQUEMENT de la nomenclature CIM-10 FR 2026
- Distingue la DESCRIPTION CLINIQUE (ce que le médecin écrit) de la LOGIQUE DE CODAGE (ce que l'ATIH impose)
- Privilégie le code le plus SPÉCIFIQUE disponible (4e ou 5e caractère)
- Vérifie les notes d'inclusion/exclusion de chaque code candidat
- Un DAS doit avoir mobilisé des ressources supplémentaires pendant le séjour
- EXCLUSION SYMPTÔME : Si le diagnostic est un symptôme (R00-R99) et qu'un diagnostic précis (Chapitres I-XIV, A00-N99) expliquant ce symptôme est présent, le symptôme ne doit PAS être codé

DIAGNOSTIC À CODER : "{texte}"
TYPE : DAS (diagnostic associé significatif)

CONTEXTE CLINIQUE :
{contexte}

SOURCES CIM-10 :
{sources}

Réponds UNIQUEMENT avec un objet JSON au format suivant, sans aucun texte avant ou après :
{{
  "analyse_clinique": "que signifie ce diagnostic sur le plan médical",
  "codes_candidats": "quels codes CIM-10 des sources sont compatibles",
  "discrimination": "pourquoi choisir ce code plutôt qu'un autre (inclusions/exclusions, spécificité)",
  "regle_pmsi": "conformité aux règles PMSI pour un DAS (guide méthodologique)",
  "code": "X99.9",
  "confidence": "high ou medium ou low",
  "justification": "explication courte en français"
}}"""


# V2 : Template CPAM contre-argumentation (3e variante)
PROMPT_TEMPLATE_CPAM = """Tu es médecin DIM. La CPAM conteste le codage d'un diagnostic. Argumente.

DIAGNOSTIC CONTESTÉ : "{texte}"
TYPE : {type_label}
CODE ACTUEL : {code}

MOTIF DE CONTESTATION CPAM :
{motif_cpam}

SOURCES CIM-10 :
{sources}

Réponds UNIQUEMENT avec un objet JSON :
{{
  "analyse_clinique": "que signifie ce diagnostic sur le plan médical",
  "codes_candidats": "quels codes CIM-10 sont compatibles",
  "discrimination": "pourquoi ce code est correct (inclusions/exclusions, spécificité)",
  "regle_pmsi": "conformité aux règles PMSI et au guide méthodologique",
  "code": "{code}",
  "confidence": "high ou medium ou low",
  "justification": "argumentation structurée pour répondre à la contestation CPAM"
}}"""

CPAM_MOTIFS = [
    "Le code {code} ne semble pas justifié par les éléments du dossier médical.",
    "La CPAM propose de recoder en {alt_code} ({alt_label}). Justifiez le maintien du code actuel.",
    "Le DAS {code} n'a pas mobilisé de ressources supplémentaires pendant le séjour.",
    "Ce diagnostic est un symptôme déjà couvert par le DP. Il ne devrait pas être codé séparément.",
    "Le niveau de sévérité CMA associé à {code} ne semble pas justifié cliniquement.",
]


def build_fake_source(code: str, cim10_dict: dict[str, str]) -> str:
    """Génère un extrait de source CIM-10 simulé pour un code donné."""
    label = cim10_dict.get(code, "")
    if not label:
        return ""

    # Trouver des codes voisins (même catégorie à 3 car.)
    prefix = code[:3]
    neighbors = []
    for c, l in cim10_dict.items():
        if c[:3] == prefix and c != code:
            neighbors.append((c, l))
    neighbors.sort()
    neighbors = neighbors[:4]

    lines = [f"--- Source 1: CIM-10 FR 2026 (code: {code}) ---"]
    lines.append(f"{code} {label}")
    for c, l in neighbors:
        lines.append(f"{c} {l}")
    lines.append("")

    return "\n".join(lines)


def build_response_json(entry: dict) -> str:
    """Reconstruit la réponse JSON structurée depuis une entrée du cache."""
    # Extraire les sections du raisonnement
    rais = entry.get("raisonnement", "")
    analyse = ""
    candidats = ""
    discrim = ""
    regle = ""

    if "ANALYSE CLINIQUE" in rais:
        parts = rais.split("\n\n")
        for part in parts:
            p = part.strip()
            if p.startswith("ANALYSE CLINIQUE"):
                analyse = p.replace("ANALYSE CLINIQUE :\n", "").replace("ANALYSE CLINIQUE :", "").strip()
            elif p.startswith("CODES CANDIDATS"):
                candidats = p.replace("CODES CANDIDATS :\n", "").replace("CODES CANDIDATS :", "").strip()
            elif p.startswith("DISCRIMINATION"):
                discrim = p.replace("DISCRIMINATION :\n", "").replace("DISCRIMINATION :", "").strip()
            elif p.startswith("REGLE PMSI") or p.startswith("RÈGLE PMSI"):
                regle = p.split(":\n", 1)[-1].strip() if ":\n" in p else p.split(":", 1)[-1].strip()

    resp = {
        "analyse_clinique": analyse or "Diagnostic médical nécessitant un codage CIM-10 spécifique.",
        "codes_candidats": candidats or f"[{entry.get('code', '?')}]",
        "discrimination": discrim or entry.get("justification", ""),
        "regle_pmsi": regle or "Code conforme aux règles PMSI.",
        "code": entry.get("code", "?"),
        "confidence": entry.get("confidence", "medium"),
        "justification": entry.get("justification", ""),
    }

    return json.dumps(resp, ensure_ascii=False)


def generate_contexte_samples() -> list[str]:
    """Génère des contextes patients variés."""
    return [
        "- Patient : Homme, 72 ans, IMC 28.5\n- Durée séjour : 7 jours\n- Biologie : CRP 145 [N: 0-5] (↑), Créatinine 89 [N: 50-120]",
        "- Patient : Femme, 58 ans, IMC 24.1\n- Durée séjour : 3 jours\n- Biologie : Hémoglobine 10.2 [N: 12-17] (↑), Plaquettes 180 [N: 150-400]",
        "- Patient : Homme, 85 ans, IMC 22.0\n- Durée séjour : 12 jours\n- Antécédents : HTA, Diabète de type 2, BPCO\n- Biologie : CRP 89 [N: 0-5] (↑), Leucocytes 14.2 [N: 4-10] (↑)",
        "- Patient : Femme, 45 ans\n- Durée séjour : 2 jours",
        "- Patient : Homme, 67 ans, IMC 31.2\n- Durée séjour : 5 jours\n- Biologie : ASAT 85 [N: 0-40] (↑), ALAT 120 [N: 0-40] (↑), GGT 210 [N: 0-60] (↑)",
        "- Patient : Femme, 30 ans\n- Durée séjour : 4 jours\n- Biologie : CRP 25 [N: 0-5] (↑), Leucocytes 12.5 [N: 4-10] (↑)",
        "- Patient : Homme, 78 ans, IMC 20.1\n- Durée séjour : 14 jours\n- Antécédents : Insuffisance cardiaque, FA\n- Complications : Infection urinaire",
        "Non précisé",
    ]


def _parse_cache_key(key):
    """Extraire le type (dp/das) et le texte depuis la clé du cache."""
    if key.startswith("das_llm::das_extract::"):
        parts = key.split("::", 3)
        texte = parts[3] if len(parts) > 3 else parts[-1]
        return "das", texte.strip()
    if "::" in key:
        diag_type, texte = key.split("::", 1)
        return diag_type.strip(), texte.strip()
    return "das", key.strip()


def _find_alt_code(code: str, cim10_dict: dict[str, str]) -> tuple[str, str]:
    """Trouve un code alternatif (même catégorie) pour les motifs CPAM."""
    prefix = code[:3]
    for c, label in cim10_dict.items():
        if c[:3] == prefix and c != code and len(c) >= 4:
            return c, label
    # Fallback : code .9 (SAI)
    sai = f"{prefix}.9"
    if sai in cim10_dict and sai != code:
        return sai, cim10_dict[sai]
    return "Z03.9", "Observation pour suspicion non précisée"


def main():
    print("Chargement du cache Ollama gemma3:12b (toutes sources)...")
    entries = load_cache_entries()
    print(f"  Total fusionné : {len(entries)} entrées")

    print("Chargement du dictionnaire CIM-10...")
    cim10_dict = load_cim10_dict()
    print(f"  {len(cim10_dict)} codes")

    contextes = generate_contexte_samples()
    examples = []
    n_short = 0
    n_long = 0
    n_cpam = 0

    for key, entry in entries.items():
        code = entry.get("code", "")
        if not code or code == "?" or len(code) < 3:
            continue

        diag_type, texte = _parse_cache_key(key)
        if diag_type not in ("dp", "das"):
            continue
        if not texte or len(texte) < 3:
            continue

        # Vérifier que le raisonnement est substantiel
        raisonnement = entry.get("raisonnement", "")
        if len(raisonnement) < 50:
            continue

        response_json = build_response_json(entry)

        # V2 : TOUJOURS générer version courte + version longue

        # 1) Version courte (sans sources ni contexte)
        if diag_type == "dp":
            user_content = PROMPT_TEMPLATE_DP.format(texte=texte)
        else:
            user_content = PROMPT_TEMPLATE_DAS.format(texte=texte)

        examples.append({
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_content},
                {"role": "assistant", "content": response_json},
            ]
        })
        n_short += 1

        # 2) Version longue (avec sources CIM-10 et contexte simulé)
        source_text = build_fake_source(code, cim10_dict)
        if source_text:
            ctx = random.choice(contextes)
            if diag_type == "dp":
                user_long = PROMPT_TEMPLATE_DP_LONG.format(
                    texte=texte, contexte=ctx, sources=source_text
                )
            else:
                user_long = PROMPT_TEMPLATE_DAS_LONG.format(
                    texte=texte, contexte=ctx, sources=source_text
                )

            examples.append({
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_long},
                    {"role": "assistant", "content": response_json},
                ]
            })
            n_long += 1

        # 3) V2 : Version CPAM (contre-argumentation) — 20% des entrées DAS high confidence
        if (diag_type == "das"
                and entry.get("confidence") == "high"
                and source_text
                and random.random() < 0.20):
            type_label = "DAS (diagnostic associé significatif)"
            alt_code, alt_label = _find_alt_code(code, cim10_dict)
            motif = random.choice(CPAM_MOTIFS).format(
                code=code, alt_code=alt_code, alt_label=alt_label
            )
            user_cpam = PROMPT_TEMPLATE_CPAM.format(
                texte=texte, type_label=type_label, code=code,
                motif_cpam=motif, sources=source_text,
            )
            examples.append({
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": user_cpam},
                    {"role": "assistant", "content": response_json},
                ]
            })
            n_cpam += 1

    # Mélanger
    random.shuffle(examples)

    # Écrire
    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    with open(OUTPUT, "w") as f:
        for ex in examples:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    # Stats
    token_counts = []
    for ex in examples:
        text = " ".join(m["content"] for m in ex["messages"])
        token_counts.append(int(len(text.split()) * 1.3))

    print(f"\n{'='*50}")
    print(f"Exemples pipeline générés : {len(examples)}")
    print(f"  Courts : {n_short}, Longs : {n_long}, CPAM : {n_cpam}")
    print(f"  → {OUTPUT}")
    print(f"  Tokens : moy={sum(token_counts)//len(token_counts)}, "
          f"max={max(token_counts)}, min={min(token_counts)}")
    print(f"  Taille : {OUTPUT.stat().st_size / 1024 / 1024:.1f} Mo")


if __name__ == "__main__":
    main()