export inference-ready

inference.py

@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+# llmTalk_ids_v8_hf.py
+# ============================================================
+# INFERENCE EN IDS UNIQUEMENT (vocab=8):
+# 0/1 bits + 6 specials: BOS EOS BOI EOI BOR EOR
+#
+# Deux modes de prompt:
+# - --prompt_ids : string de chiffres (ex: "240000001540000015") (digits only, 0..7) (peut être "")
+# - --prompt_int : string "int,int" -> génère: BOS t0 t1 BOI int1(10b) EOI BOI int2(10b) EOI
+#
+# Option:
+# - --print_int : extrait le premier bloc BOR ... EOR (bits variables) dans la séquence complète
+#                et affiche sa valeur décimale (binaire -> int).
+#                (min_bits=10 par défaut pour coller à tes entrées 10 bits, mais la réponse peut dépasser)
+# ============================================================
+
+import sys
+import argparse
+import random
+from collections import Counter
+from typing import List, Dict, Tuple, Any, Optional
+
+import torch
+from transformers import AutoModelForCausalLM
+
+# ----------------------------
+# Special tokens (vocab=8)
+# ----------------------------
+TOK_BOS = 2
+TOK_EOS = 3
+TOK_BOI = 4
+TOK_EOI = 5
+TOK_BOR = 6
+TOK_EOR = 7
+
+TOK_NAMES = {
+    0: "0",
+    1: "1",
+    TOK_BOS: "BOS",
+    TOK_EOS: "EOS",
+    TOK_BOI: "BOI",
+    TOK_EOI: "EOI",
+    TOK_BOR: "BOR",
+    TOK_EOR: "EOR",
+}
+
+# ------------------------------------------------------------
+# Task header bits for --prompt_int (t0, t1)
+# ------------------------------------------------------------
+# Tu as demandé "BOS t0 t1 ...", sans préciser t0/t1.
+# Ici je mets un défaut neutre: 0,0 (modifiable si tu veux).
+PROMPT_INT_T0 = 0
+PROMPT_INT_T1 = 0
+
+# ----------------------------
+# Logits modifiers
+# ----------------------------
+def apply_repetition_penalty_(logits: torch.Tensor, token_ids: List[int], penalty: float) -> None:
+    if penalty is None or penalty == 1.0 or penalty <= 0:
+        return
+    for t in set(token_ids):
+        val = logits[0, t]
+        logits[0, t] = val * penalty if val < 0 else val / penalty
+
+def apply_encoder_repetition_penalty_(logits: torch.Tensor, prompt_token_ids: List[int], penalty: float) -> None:
+    if penalty is None or penalty == 1.0 or penalty <= 0:
+        return
+    for t in set(prompt_token_ids):
+        val = logits[0, t]
+        logits[0, t] = val / penalty if val < 0 else val * penalty
+
+def apply_presence_frequency_penalties_(
+    logits: torch.Tensor,
+    token_ids: List[int],
+    presence_penalty: float,
+    frequency_penalty: float,
+) -> None:
+    counts = Counter(token_ids)
+
+    if presence_penalty:
+        for t in counts:
+            logits[0, t] -= presence_penalty
+
+    if frequency_penalty:
+        for t, c in counts.items():
+            logits[0, t] -= frequency_penalty * c
+
+def get_banned_tokens_no_repeat_ngram(seq: List[int], n: int) -> set:
+    if n <= 0 or len(seq) < n - 1:
+        return set()
+
+    prefix_len = n - 1
+    ngrams: Dict[Tuple[int, ...], set] = {}
+    for i in range(len(seq) - n + 1):
+        prefix = tuple(seq[i:i + prefix_len])
+        nxt = seq[i + prefix_len]
+        ngrams.setdefault(prefix, set()).add(nxt)
+
+    return ngrams.get(tuple(seq[-prefix_len:]), set())
+
+def mask_banned_tokens_(logits: torch.Tensor, banned: set) -> None:
+    if banned:
+        logits[0, list(banned)] = float("-inf")
+
+# ----------------------------
+# Helpers: prompt parsing + pretty print
+# ----------------------------
+def parse_prompt_ids_str(s: str, vocab_size: int = 8) -> List[int]:
+    s = "" if s is None else str(s)
+    s = s.strip()
+    if s == "":
+        return []
+
+    if not s.isdigit():
+        raise ValueError("prompt_ids doit contenir uniquement des chiffres (0..7), sans espaces.")
+
+    ids: List[int] = []
+    for ch in s:
+        t = ord(ch) - ord("0")
+        if t < 0 or t >= vocab_size:
+            raise ValueError(f"token id hors vocab: {t} (vocab_size={vocab_size})")
+        ids.append(t)
+    return ids
+
+def format_ids_readable(ids: List[int]) -> str:
+    out: List[str] = []
+    for t in ids:
+        out.append(TOK_NAMES.get(int(t), str(int(t))))
+    return " ".join(out)
+
+def format_ids_compact(ids: List[int]) -> str:
+    s: List[str] = []
+    for t in ids:
+        ti = int(t)
+        if ti in (0, 1):
+            if s and (s[-1] and s[-1][-1] in ("0", "1")):
+                s[-1] = s[-1] + str(ti)
+            else:
+                s.append(str(ti))
+        else:
+            s.append(TOK_NAMES.get(ti, str(ti)))
+    return " ".join(s)
+
+# ----------------------------
+# --prompt_int builder
+# ----------------------------
+def int_to_10bits_tokens(x: int) -> List[int]:
+    if x < 0 or x > 1023:
+        raise ValueError(f"int hors range pour 10 bits: {x} (attendu 0..1023)")
+    b = format(int(x), "010b")  # MSB -> LSB
+    return [0 if ch == "0" else 1 for ch in b]
+
+def parse_prompt_int_str(s: str) -> Tuple[int, int]:
+    s = "" if s is None else str(s)
+    s = s.strip()
+    if s == "":
+        raise ValueError("--prompt_int vide. Attendu: \"int,int\"")
+
+    parts = s.split(",")
+    if len(parts) != 2:
+        raise ValueError(f"--prompt_int invalide: {s!r}. Attendu: \"int,int\"")
+
+    try:
+        a = int(parts[0].strip())
+        b = int(parts[1].strip())
+    except Exception:
+        raise ValueError(f"--prompt_int invalide: {s!r}. Les deux valeurs doivent être des int.")
+
+    return a, b
+
+def build_prompt_from_ints(int1: int, int2: int) -> List[int]:
+    seq: List[int] = []
+    seq.append(TOK_BOS)
+    seq.append(int(PROMPT_INT_T0))
+    seq.append(int(PROMPT_INT_T1))
+
+    seq.append(TOK_BOI)
+    seq.extend(int_to_10bits_tokens(int1))
+    seq.append(TOK_EOI)
+
+    seq.append(TOK_BOI)
+    seq.extend(int_to_10bits_tokens(int2))
+    seq.append(TOK_EOI)
+
+    return seq
+
+# ----------------------------
+# --print_int extractor (BOR ... EOR, bits variables)
+# ----------------------------
+def extract_first_bor_eor_bits(ids: List[int], min_bits: int = 1) -> Optional[Tuple[List[int], int, int]]:
+    try:
+        i = ids.index(TOK_BOR)
+    except ValueError:
+        return None
+
+    bits: List[int] = []
+    j = i + 1
+    while j < len(ids):
+        t = int(ids[j])
+        if t == TOK_EOR:
+            break
+        if t in (0, 1):
+            bits.append(t)
+        j += 1
+
+    if len(bits) < int(min_bits):
+        return None
+
+    val = 0
+    for b in bits:
+        val = (val << 1) | int(b)
+
+    return bits, val, i
+
+# ----------------------------
+# Main
+# ----------------------------
+def main() -> None:
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--repo", type=str, required=True, help='HF repo id ou path local (ex: "PhysiQuanty/xxx")')
+    parser.add_argument("--revision", type=str, default=None, help="HF revision/branch/tag/commit (optionnel)")
+
+    g = parser.add_mutually_exclusive_group(required=False)
+    g.add_argument("--prompt_ids", type=str, default=None, help='Ex: "240000001540000015" (digits only 0..7) or ""')
+    g.add_argument("--prompt_int", type=str, default=None, help='Ex: "12,900" -> BOS t0 t1 BOI 10b EOI BOI 10b EOI')
+
+    parser.add_argument("--print_int", action="store_true", help="Affiche le 1er bloc BOR..EOR (bits) en int")
+
+    parser.add_argument("--max_new_tokens", type=int, default=40)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top_k", type=int, default=50)
+
+    parser.add_argument("--repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--presence_penalty", type=float, default=0.0)
+    parser.add_argument("--frequency_penalty", type=float, default=0.0)
+    parser.add_argument("--encoder_repetition_penalty", type=float, default=1.0)
+    parser.add_argument("--no_repeat_ngram_size", type=int, default=0)
+
+    parser.add_argument("--seed", type=int, default=-1)
+    parser.add_argument("--device", type=str, default="cuda", choices=["cpu", "cuda"])
+
+    parser.add_argument("--stream_ids", action="store_true", help="Stream les IDS générés au fil de l'eau")
+    parser.add_argument("--print_prompt_readable", action="store_true", help="Affiche prompt en tokens lisibles")
+    parser.add_argument("--print_final_readable", action="store_true", help="Affiche sortie finale en tokens lisibles")
+    parser.add_argument("--stop_on_eos", action="store_true", help="Stop dès que EOS(3) est généré")
+
+    args = parser.parse_args()
+
+    seed = args.seed if args.seed >= 0 else random.randint(0, 2**31 - 1)
+    print(f"[Seed] {seed}", flush=True)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    device = torch.device("cuda" if (args.device == "cuda" and torch.cuda.is_available()) else "cpu")
+    print(f"[Device] {device}", flush=True)
+
+    torch_dtype = torch.float16 if device.type == "cuda" else torch.float32
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.repo,
+        revision=args.revision,
+        trust_remote_code=True,
+        torch_dtype=torch_dtype,
+        low_cpu_mem_usage=True,
+    )
+    model.to(device)
+    model.eval()
+
+    vocab_size_cfg = int(getattr(model.config, "vocab_size", -1))
+    print(f"[Model] loaded from {args.repo} | vocab_size={vocab_size_cfg}", flush=True)
+    if vocab_size_cfg != 8:
+        print(f"[Warn] vocab_size={vocab_size_cfg} (attendu 8).", flush=True)
+
+    # ---- build prompt ids from either --prompt_int or --prompt_ids (or default "")
+    if args.prompt_int is not None:
+        int1, int2 = parse_prompt_int_str(args.prompt_int)
+        prompt_ids = build_prompt_from_ints(int1, int2)
+        prompt_origin = f'prompt_int="{args.prompt_int}" (t0,t1={PROMPT_INT_T0},{PROMPT_INT_T1})'
+    else:
+        s = "" if args.prompt_ids is None else args.prompt_ids
+        prompt_ids = parse_prompt_ids_str(s, vocab_size=8)
+        prompt_origin = 'prompt_ids' if args.prompt_ids is not None else 'prompt_ids="" (default)'
+
+    print(f"[Prompt Origin] {prompt_origin}", flush=True)
+
+    if args.print_prompt_readable:
+        print(f"[Prompt IDs] {prompt_ids}", flush=True)
+        print(f"[Prompt readable] {format_ids_readable(prompt_ids)}", flush=True)
+        print(f"[Prompt compact]  {format_ids_compact(prompt_ids)}", flush=True)
+    else:
+        if len(prompt_ids) == 0:
+            print("[Prompt IDs] len=0 (prompt nul)", flush=True)
+        else:
+            print(f"[Prompt IDs] len={len(prompt_ids)} first32={prompt_ids[:32]}", flush=True)
+
+    seeded_with_bos = False
+    if len(prompt_ids) == 0:
+        tokens = torch.tensor([TOK_BOS], device=device, dtype=torch.long).unsqueeze(0)
+        seeded_with_bos = True
+    else:
+        tokens = torch.tensor(prompt_ids, device=device, dtype=torch.long).unsqueeze(0)
+
+    generated_raw: List[int] = []
+
+    if args.stream_ids:
+        sys.stdout.write("[Stream IDS] ")
+        sys.stdout.flush()
+
+    with torch.no_grad():
+        for _ in range(int(args.max_new_tokens)):
+            out = model(input_ids=tokens)
+            logits = out.logits[:, -1, :]  # (1, vocab)
+
+            logits_work = logits.clone()
+            full_seq = tokens[0].tolist()
+
+            apply_encoder_repetition_penalty_(logits_work, prompt_ids, float(args.encoder_repetition_penalty))
+            apply_repetition_penalty_(logits_work, full_seq, float(args.repetition_penalty))
+            apply_presence_frequency_penalties_(
+                logits_work,
+                full_seq,
+                float(args.presence_penalty),
+                float(args.frequency_penalty),
+            )
+
+            if int(args.no_repeat_ngram_size) > 0:
+                banned = get_banned_tokens_no_repeat_ngram(full_seq, int(args.no_repeat_ngram_size))
+                mask_banned_tokens_(logits_work, banned)
+
+            logits_work /= max(float(args.temperature), 1e-6)
+
+            if 0 < int(args.top_k) < logits_work.size(-1):
+                v, _ = torch.topk(logits_work, int(args.top_k))
+                logits_work[logits_work < v[:, [-1]]] = float("-inf")
+
+            probs = torch.softmax(logits_work, dim=-1)
+            next_token = torch.multinomial(probs, 1)  # (1,1)
+            tok_id = int(next_token.item())
+            generated_raw.append(tok_id)
+
+            if args.stream_ids:
+                sys.stdout.write(str(tok_id))
+                sys.stdout.flush()
+
+            tokens = torch.cat([tokens, next_token], dim=1)
+
+            if args.stop_on_eos and tok_id == TOK_EOS:
+                break
+
+    if args.stream_ids:
+        sys.stdout.write("\n")
+        sys.stdout.flush()
+
+    if seeded_with_bos:
+        print("\n[Prompt] prompt nul -> seed interne BOS(2) utilisé uniquement pour init logits", flush=True)
+
+    print("\n[Generated RAW IDS]", flush=True)
+    print(generated_raw, flush=True)
+
+    print("\n[Generated RAW IDS (as digits)]", flush=True)
+    print("".join(str(x) for x in generated_raw), flush=True)
+
+    if args.print_final_readable or args.print_int:
+        full = prompt_ids + generated_raw
+
+    if args.print_final_readable:
+        print("\n[Full sequence readable]", flush=True)
+        print(format_ids_readable(full), flush=True)
+        print("\n[Full sequence compact]", flush=True)
+        print(format_ids_compact(full), flush=True)
+
+    if args.print_int:
+        got = extract_first_bor_eor_bits(full, min_bits=10)
+        if got is None:
+            print("\n[PrintInt] Aucun bloc BOR..EOR valide trouvé.", flush=True)
+        else:
+            bits, val, pos = got
+            bits_str = "".join(str(b) for b in bits)
+            print("\n[PrintInt] First BOR..EOR", flush=True)
+            print(f"[PrintInt] pos={pos} nbits={len(bits)} bits={bits_str} int={val}", flush=True)
+
+if __name__ == "__main__":
+    main()