Upload distill/generate.py with huggingface_hub

distill/generate.py

@@ -0,0 +1,280 @@
+"""Étape 1 — collecte des réponses de Kimi-K3 pour distiller vers Kimi-Linear.
+
+L'API officielle ne rend pas les logits : la distillation est donc au niveau
+séquence (le socle apprend à imiter les sorties de K3), pas au niveau
+distribution. C'est moins efficace en tokens qu'une KD sur logits, mais c'est
+la seule voie disponible sans poids K3 locaux -- et c'est la recette déjà
+validée dans Kimi-K3-L4-Distill-9B-LoRA.
+
+Le budget est un plafond dur, vérifié avant chaque appel : l'endpoint facture
+~15 $/1M de tokens de sortie, et une boucle non bornée y passerait des crédits
+sans prévenir.
+
+    python generate.py --budget 2.0 --out data/k3.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import pathlib
+import random
+import sys
+import time
+import urllib.error
+import urllib.request
+
+ENDPOINT = "https://api.runpod.ai/v2/moonshot-kimi/runsync"
+PRICE_OUT = 15.0 / 1_000_000     # $/token de sortie, tarif affiche
+PRICE_IN = 3.0 / 1_000_000
+
+
+def api_key() -> str:
+    """Le jeton n'est jamais journalise ni affiche."""
+    k = os.environ.get("RUNPOD_API_KEY")
+    if k:
+        return k.strip()
+    for p in (r"G:\AI\runpod\.env", ".env"):
+        f = pathlib.Path(p)
+        if f.exists():
+            for line in f.read_text(encoding="utf-8", errors="ignore").splitlines():
+                if line.startswith("RUNPOD_API_KEY="):
+                    v = line.split("=", 1)[1].strip().strip('"').strip("'")
+                    if v:
+                        return v
+    sys.exit("RUNPOD_API_KEY introuvable (env ou G:\\AI\\runpod\\.env)")
+
+
+# ---------------------------------------------------------------------- prompts
+# Cible les écarts mesurés entre Kimi-Linear et K3, pas des tâches génériques :
+# le français imprécis ("gate" traduit en "gâteau"), l'usage d'outils, et
+# l'édition de code, qui est le régime réel d'un agent.
+
+CODE = [
+    "Écris une fonction Python qui parse un fichier .env en dictionnaire, en gérant les commentaires, les guillemets et les lignes vides. Code seulement.",
+    "Refactorise ce code pour éliminer la duplication :\n\n```python\ndef get_user(id):\n    r = requests.get(f'/users/{id}')\n    if r.status_code != 200: raise Exception('fail')\n    return r.json()\n\ndef get_post(id):\n    r = requests.get(f'/posts/{id}')\n    if r.status_code != 200: raise Exception('fail')\n    return r.json()\n```",
+    "Explique pourquoi ce code a une condition de course et corrige-le :\n\n```python\ncounter = 0\ndef worker():\n    global counter\n    for _ in range(100000):\n        counter += 1\n```",
+    "Écris un décorateur Python qui réessaie une fonction avec un backoff exponentiel et un jitter. Type hints complets.",
+    "Convertis cette boucle en compréhension et explique si c'est vraiment plus lisible :\n\n```python\nout = []\nfor x in items:\n    if x.active:\n        out.append(x.name.upper())\n```",
+    "Trouve le bug :\n\n```python\ndef binary_search(a, t):\n    lo, hi = 0, len(a)\n    while lo < hi:\n        mid = (lo + hi) // 2\n        if a[mid] < t: lo = mid\n        elif a[mid] > t: hi = mid\n        else: return mid\n    return -1\n```",
+    "Écris un contexte manager asynchrone Python qui limite la concurrence à N tâches.",
+    "Implémente un LRU cache en Python sans utiliser functools, avec O(1) en lecture et écriture.",
+]
+
+FRENCH = [
+    "Explique en trois phrases précises ce qu'est le mécanisme de gating dans un modèle mixture-of-experts. Emploie la terminologie française correcte.",
+    "Quelle est la différence entre la mémoïsation et la programmation dynamique ? Réponds en français technique rigoureux.",
+    "Explique ce qu'est une condition de concurrence et donne un exemple concret en français.",
+    "Décris le fonctionnement d'un cache de préfixe dans un moteur d'inférence, en français, sans anglicismes inutiles.",
+    "Explique la différence entre latence et débit pour un serveur d'inférence, avec un exemple chiffré.",
+    "Qu'est-ce que la quantification d'un modèle de langage ? Explique les compromis en français précis.",
+]
+
+# Les outils sont passes dans le champ `tools` de la requete, pas decrits en
+# prose : sollicite ainsi, K3 repond a juste titre "aucun outil ne m'est
+# reellement fourni", et l'exemple collecte est inutilisable pour apprendre
+# l'usage d'outils.
+TOOL_SCHEMA = [
+    {"type": "function", "function": {"name": "read_file", "description": "Lit un fichier",
+     "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}}},
+    {"type": "function", "function": {"name": "write_file", "description": "Ecrit un fichier",
+     "parameters": {"type": "object", "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, "required": ["path", "content"]}}},
+    {"type": "function", "function": {"name": "grep", "description": "Cherche un motif",
+     "parameters": {"type": "object", "properties": {"pattern": {"type": "string"}, "path": {"type": "string"}}, "required": ["pattern"]}}},
+    {"type": "function", "function": {"name": "list_dir", "description": "Liste un dossier",
+     "parameters": {"type": "object", "properties": {"path": {"type": "string"}}, "required": ["path"]}}},
+    {"type": "function", "function": {"name": "run_command", "description": "Execute une commande shell",
+     "parameters": {"type": "object", "properties": {"cmd": {"type": "string"}}, "required": ["cmd"]}}},
+]
+
+TOOLS = [
+    "Lis /etc/config.yaml, remplace le port 8080 par 9090, et reecris le fichier.",
+    "Trouve tous les fichiers Python de plus de 500 lignes dans ce depot.",
+    "Trouve ou la variable MAX_RETRIES est definie, puis montre son contexte.",
+    "Determine quel framework de test utilise ce projet.",
+]
+
+REASONING = [
+    "Un serveur traite 1200 requêtes/s avec 8 workers. On passe à 12 workers et le débit monte à 1500 req/s. Que peut-on en déduire sur le goulot d'étranglement ?",
+    "Un modèle lit 4,7 Go de poids par token et la carte offre 696 Go/s. Quel est le plafond théorique en tokens/s ? Comment le décodage spéculatif le dépasse-t-il ?",
+    "Pourquoi un cache KV quantifié en 8 bits peut-il être refusé par un moteur alors que les poids en 4 bits sont acceptés ?",
+]
+
+
+def build_prompts(n: int, seed: int = 0) -> list[str]:
+    """Mélange déterministe : deux exécutions avec la même graine donnent le
+    même jeu, ce qui rend une reprise après interruption reproductible."""
+    pool = CODE * 3 + FRENCH * 3 + TOOLS * 2 + REASONING * 2
+    rng = random.Random(seed)
+    rng.shuffle(pool)
+    return (pool * (n // len(pool) + 1))[:n]
+
+
+def _req(url: str, key: str, body: bytes | None = None, timeout: int = 300):
+    req = urllib.request.Request(
+        url, data=body,
+        headers={"Authorization": f"Bearer {key}",
+                 "Content-Type": "application/json",
+                 # Cloudflare, devant l'API Runpod, rejette la signature par
+                 # defaut d'urllib avec un 403 "error code: 1010".
+                 "User-Agent": "Mozilla/5.0 (distill)"})
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return json.loads(r.read())
+
+
+def call_k3(prompt: str, key: str, max_tokens: int, timeout: int = 300,
+            tools: list | None = None) -> dict | None:
+    """Soumet puis, si necessaire, suit la tache jusqu'a son terme.
+
+    `runsync` ne rend le resultat que si la generation tient dans son attente
+    interne ; au-dela il renvoie un identifiant et un statut IN_QUEUE ou
+    IN_PROGRESS. Ne pas le suivre donnait une reponse sans `output`, donc vide
+    et sans cout -- un symptome qui ressemblait a tort a un probleme de modele.
+    """
+    body = json.dumps({"input": {
+        "model": "kimi-k3",
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        # K3 refuse toute temperature autre que 1 :
+        #   "invalid temperature: only 1 is allowed for this model"
+        # L'erreur n'apparaissait pas tant que la tache n'etait pas suivie.
+        "temperature": 1,
+        **({"tools": tools, "tool_choice": "auto"} if tools else {}),
+    }}).encode()
+    try:
+        r = _req("https://api.runpod.ai/v2/moonshot-kimi/runsync", key, body, timeout)
+    except urllib.error.HTTPError as e:
+        print(f"  HTTP {e.code}: {e.read()[:200].decode('utf-8', 'ignore')}", flush=True)
+        return None
+    except Exception as e:
+        print(f"  {type(e).__name__}: {e}", flush=True)
+        return None
+
+    if r.get("output") is not None:
+        return r
+
+    job = r.get("id")
+    status = r.get("status")
+    if not job:
+        print(f"  reponse sans output ni id : {str(r)[:160]}", flush=True)
+        return None
+
+    deadline = time.time() + timeout
+    while time.time() < deadline:
+        time.sleep(3)
+        try:
+            r = _req(f"https://api.runpod.ai/v2/moonshot-kimi/status/{job}", key)
+        except Exception as e:
+            print(f"  suivi : {type(e).__name__}", flush=True)
+            continue
+        status = r.get("status")
+        if status == "COMPLETED":
+            return r
+        if status in ("FAILED", "CANCELLED", "TIMED_OUT"):
+            print(f"  tache {status}: {str(r.get('error'))[:160]}", flush=True)
+            return None
+    print(f"  tache toujours {status} apres {timeout} s", flush=True)
+    return None
+
+
+def extract(resp: dict) -> tuple[str, str, dict, float, list]:
+    """Renvoie (contenu, trace de raisonnement, usage, coût réel).
+
+    K3 expose son raisonnement dans `reasoning_content`. On le conserve à part :
+    l'entraîner ou non est une décision à prendre en aval, pas ici.
+    """
+    out = resp.get("output")
+    if isinstance(out, list) and out:
+        out = out[0]
+    res = (out or {}).get("result") or {}
+    ch = (res.get("choices") or [{}])[0]
+    msg = ch.get("message") or {}
+    usage = res.get("usage") or {}
+    cost = (out or {}).get("cost")
+    if cost is None:
+        cost = (usage.get("completion_tokens", 0) * PRICE_OUT
+                + usage.get("prompt_tokens", 0) * PRICE_IN)
+    return (msg.get("content") or "", msg.get("reasoning_content") or "",
+            usage, float(cost), msg.get("tool_calls") or [])
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--budget", type=float, default=1.0, help="plafond dur en dollars")
+    ap.add_argument("--out", default="data/k3.jsonl")
+    # K3 raisonne avant de repondre : un budget trop court ne produit que de la
+    # reflexion. 2000 laisse la place a la trace ET a la reponse.
+    ap.add_argument("--max-tokens", type=int, default=2000)
+    ap.add_argument("--n", type=int, default=400)
+    ap.add_argument("--seed", type=int, default=0)
+    a = ap.parse_args()
+
+    key = api_key()
+    path = pathlib.Path(a.out)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Reprise : on saute ce qui est déjà collecté plutôt que de le repayer.
+    done = set()
+    if path.exists():
+        for line in path.open(encoding="utf-8"):
+            try:
+                done.add(json.loads(line)["prompt"])
+            except Exception:
+                pass
+        print(f"reprise : {len(done)} exemples deja collectes")
+
+    prompts = [p for p in build_prompts(a.n, a.seed) if p not in done]
+    spent = 0.0
+    kept = 0
+    t0 = time.time()
+
+    with path.open("a", encoding="utf-8") as f:
+        for i, p in enumerate(prompts, 1):
+            # Verification AVANT l'appel : depasser le plafond puis s'en
+            # apercevoir ne sert a rien, l'argent est deja parti.
+            if spent >= a.budget:
+                print(f"plafond de {a.budget} $ atteint, arret propre")
+                break
+            r = call_k3(p, key, a.max_tokens,
+                        tools=TOOL_SCHEMA if p in TOOLS else None)
+            if not r:
+                continue
+            content, reasoning, usage, cost, tool_calls = extract(r)
+            spent += cost
+            # Un tour qui n'emet que des appels d'outils est un exemple
+            # valide : c'est meme le comportement qu'on cherche a distiller.
+            if not content.strip() and not tool_calls:
+                # K3 est un modele a raisonnement : il remplit d'abord
+                # reasoning_content. Avec un max_tokens trop bas il consomme
+                # tout son budget en reflexion et n'emet jamais la reponse.
+                # Le distinguer d'une vraie reponse vide evite de chercher
+                # ailleurs.
+                why = (f"raisonnement seul ({len(reasoning)} car., "
+                       f"{usage.get('completion_tokens',0)} tokens) "
+                       f"-- augmentez --max-tokens"
+                       if reasoning.strip() else "reponse totalement vide")
+                print(f"  {i}: {why}")
+                continue
+            f.write(json.dumps({
+                "prompt": p,
+                "completion": content,
+                "reasoning": reasoning,
+                "tool_calls": tool_calls,
+                "usage": usage,
+                "cost": cost,
+            }, ensure_ascii=False) + "\n")
+            f.flush()
+            kept += 1
+            if kept % 10 == 0 or i == 1:
+                print(f"  {kept} gardes | {spent:.3f} $ / {a.budget} $ | "
+                      f"{time.time()-t0:.0f} s", flush=True)
+
+    print(f"\n{kept} exemples ecrits dans {path}")
+    print(f"cout reel : {spent:.3f} $")
+    if kept:
+        print(f"soit {spent/kept*1000:.2f} $ pour 1000 exemples")
+
+
+if __name__ == "__main__":
+    main()