Create inference.py

inference.py

@@ -0,0 +1,643 @@
+#!/usr/bin/env python3
+import argparse
+import random
+from typing import Dict, List, Optional, Set
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+IM_START = "[IM_START]"
+IM_END = "[IM_END]"
+NO_THINK = "/no think"
+
+
+# ============================================================
+# UTILS
+# ============================================================
+
+def get_dtype(name: str):
+    name = str(name).lower()
+
+    if name in {"bf16", "bfloat16"}:
+        return torch.bfloat16
+
+    if name in {"fp16", "float16", "half"}:
+        return torch.float16
+
+    if name in {"fp32", "float32", "float"}:
+        return torch.float32
+
+    raise ValueError(f"Unknown dtype: {name}")
+
+
+def set_seed(seed: int):
+    if seed is None:
+        return
+
+    if seed < 0:
+        seed = random.randint(0, 2**31 - 1)
+        print(f"[INFO] random seed: {seed}")
+
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def build_prompt(args) -> str:
+    if args.no_think:
+        return (
+            f"{IM_START}user\n"
+            f"{args.prompt} {NO_THINK}"
+            f"{IM_END}\n"
+            f"{IM_START}assistant\n"
+            "<think>\n</think>\n"
+        )
+
+    return args.prompt
+
+
+def decode(tokenizer, ids: List[int]) -> str:
+    return tokenizer.decode(ids, skip_special_tokens=False)
+
+
+def extract_completion(full_text: str, prompt: str) -> str:
+    if full_text.startswith(prompt):
+        return full_text[len(prompt):]
+
+    pos = full_text.rfind(prompt)
+    if pos != -1:
+        return full_text[pos + len(prompt):]
+
+    return full_text
+
+
+def strip_after_stop_text(text: str, stop_strings: List[str]) -> str:
+    best = None
+
+    for s in stop_strings:
+        if not s:
+            continue
+
+        pos = text.find(s)
+        if pos != -1:
+            if best is None or pos < best:
+                best = pos
+
+    if best is None:
+        return text
+
+    return text[:best]
+
+
+def build_stop_sequences(tokenizer, stop_strings: List[str]) -> List[List[int]]:
+    out = []
+
+    for s in stop_strings:
+        ids = tokenizer.encode(s, add_special_tokens=False)
+        if ids:
+            out.append(ids)
+
+    return out
+
+
+def endswith_sequence(ids: List[int], suffix: List[int]) -> bool:
+    if not suffix:
+        return False
+
+    if len(ids) < len(suffix):
+        return False
+
+    return ids[-len(suffix):] == suffix
+
+
+# ============================================================
+# SAMPLING
+# ============================================================
+
+def apply_repetition_penalty(
+    logits: torch.Tensor,
+    generated_ids: List[int],
+    penalty: float,
+) -> torch.Tensor:
+    if penalty is None or penalty == 1.0:
+        return logits
+
+    if penalty <= 0:
+        raise ValueError("--repetition-penalty must be > 0")
+
+    for tid in set(generated_ids):
+        if tid < 0 or tid >= logits.numel():
+            continue
+
+        if logits[tid] > 0:
+            logits[tid] = logits[tid] / penalty
+        else:
+            logits[tid] = logits[tid] * penalty
+
+    return logits
+
+
+def apply_frequency_presence_penalty(
+    logits: torch.Tensor,
+    generated_ids: List[int],
+    frequency_penalty: float,
+    presence_penalty: float,
+) -> torch.Tensor:
+    if not generated_ids:
+        return logits
+
+    if frequency_penalty == 0.0 and presence_penalty == 0.0:
+        return logits
+
+    counts: Dict[int, int] = {}
+
+    for tid in generated_ids:
+        counts[tid] = counts.get(tid, 0) + 1
+
+    for tid, count in counts.items():
+        if tid < 0 or tid >= logits.numel():
+            continue
+
+        if frequency_penalty:
+            logits[tid] -= frequency_penalty * count
+
+        if presence_penalty:
+            logits[tid] -= presence_penalty
+
+    return logits
+
+
+def get_banned_ngram_tokens(
+    generated_ids: List[int],
+    no_repeat_ngram_size: int,
+) -> Set[int]:
+    n = no_repeat_ngram_size
+    banned = set()
+
+    if n <= 0:
+        return banned
+
+    if len(generated_ids) + 1 < n:
+        return banned
+
+    prefix_len = n - 1
+    current_prefix = tuple(generated_ids[-prefix_len:])
+
+    ngram_map = {}
+
+    for i in range(len(generated_ids) - n + 1):
+        prefix = tuple(generated_ids[i:i + prefix_len])
+        next_token = generated_ids[i + prefix_len]
+
+        if prefix not in ngram_map:
+            ngram_map[prefix] = set()
+
+        ngram_map[prefix].add(next_token)
+
+    banned.update(ngram_map.get(current_prefix, set()))
+    return banned
+
+
+def apply_no_repeat_ngram(
+    logits: torch.Tensor,
+    generated_ids: List[int],
+    no_repeat_ngram_size: int,
+) -> torch.Tensor:
+    if no_repeat_ngram_size <= 0:
+        return logits
+
+    banned = get_banned_ngram_tokens(
+        generated_ids=generated_ids,
+        no_repeat_ngram_size=no_repeat_ngram_size,
+    )
+
+    for tid in banned:
+        if 0 <= tid < logits.numel():
+            logits[tid] = -float("inf")
+
+    return logits
+
+
+def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
+    if top_k is None or top_k <= 0:
+        return logits
+
+    top_k = min(top_k, logits.size(-1))
+    values, _ = torch.topk(logits, top_k)
+    cutoff = values[-1]
+
+    logits[logits < cutoff] = -float("inf")
+    return logits
+
+
+def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor:
+    if top_p is None or top_p >= 1.0:
+        return logits
+
+    if top_p <= 0:
+        raise ValueError("--top-p must be > 0")
+
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    sorted_probs = F.softmax(sorted_logits, dim=-1)
+    cumulative = torch.cumsum(sorted_probs, dim=-1)
+
+    remove = cumulative > top_p
+    remove[1:] = remove[:-1].clone()
+    remove[0] = False
+
+    indices_to_remove = sorted_indices[remove]
+    logits[indices_to_remove] = -float("inf")
+
+    return logits
+
+
+def apply_min_p(logits: torch.Tensor, min_p: float) -> torch.Tensor:
+    if min_p is None or min_p <= 0:
+        return logits
+
+    probs = F.softmax(logits, dim=-1)
+    max_prob = torch.max(probs)
+
+    keep = probs >= (min_p * max_prob)
+    logits[~keep] = -float("inf")
+
+    return logits
+
+
+def apply_typical_p(logits: torch.Tensor, typical_p: float) -> torch.Tensor:
+    if typical_p is None or typical_p >= 1.0:
+        return logits
+
+    if typical_p <= 0:
+        raise ValueError("--typical-p must be > 0")
+
+    probs = F.softmax(logits, dim=-1)
+    log_probs = F.log_softmax(logits, dim=-1)
+
+    entropy = -(probs * log_probs).sum()
+    shifted_scores = torch.abs((-log_probs) - entropy)
+
+    sorted_scores, sorted_indices = torch.sort(shifted_scores, descending=False)
+    sorted_probs = probs[sorted_indices]
+    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+
+    remove = cumulative_probs > typical_p
+    remove[1:] = remove[:-1].clone()
+    remove[0] = False
+
+    indices_to_remove = sorted_indices[remove]
+    logits[indices_to_remove] = -float("inf")
+
+    return logits
+
+
+def apply_bad_words(
+    logits: torch.Tensor,
+    tokenizer,
+    bad_words: List[str],
+):
+    for word in bad_words:
+        ids = tokenizer.encode(word, add_special_tokens=False)
+        if len(ids) == 1:
+            tid = ids[0]
+            if 0 <= tid < logits.numel():
+                logits[tid] = -float("inf")
+
+    return logits
+
+
+def sample_next_token(
+    logits: torch.Tensor,
+    generated_ids: List[int],
+    tokenizer,
+    args,
+) -> int:
+    logits = logits.float().clone()
+
+    logits = apply_bad_words(
+        logits=logits,
+        tokenizer=tokenizer,
+        bad_words=args.bad_words,
+    )
+
+    logits = apply_repetition_penalty(
+        logits=logits,
+        generated_ids=generated_ids,
+        penalty=args.repetition_penalty,
+    )
+
+    logits = apply_frequency_presence_penalty(
+        logits=logits,
+        generated_ids=generated_ids,
+        frequency_penalty=args.frequency_penalty,
+        presence_penalty=args.presence_penalty,
+    )
+
+    logits = apply_no_repeat_ngram(
+        logits=logits,
+        generated_ids=generated_ids,
+        no_repeat_ngram_size=args.no_repeat_ngram_size,
+    )
+
+    if args.temperature <= 0:
+        return int(torch.argmax(logits).item())
+
+    logits = logits / args.temperature
+
+    logits = apply_top_k(logits, args.top_k)
+    logits = apply_top_p(logits, args.top_p)
+    logits = apply_min_p(logits, args.min_p)
+    logits = apply_typical_p(logits, args.typical_p)
+
+    probs = F.softmax(logits, dim=-1)
+
+    if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0:
+        return int(torch.argmax(logits).item())
+
+    return int(torch.multinomial(probs, num_samples=1).item())
+
+
+# ============================================================
+# MODEL
+# ============================================================
+
+def model_forward_logits(model, input_ids: torch.Tensor):
+    out = model(input_ids=input_ids)
+
+    if hasattr(out, "logits"):
+        return out.logits
+
+    if isinstance(out, tuple):
+        return out[0]
+
+    raise RuntimeError("Impossible de récupérer logits depuis la sortie du modèle.")
+
+
+@torch.no_grad()
+def generate_manual(
+    model,
+    tokenizer,
+    input_ids: torch.Tensor,
+    args,
+) -> torch.Tensor:
+    idx = input_ids
+    generated_after_prompt: List[int] = []
+
+    stop_sequences = build_stop_sequences(
+        tokenizer,
+        stop_strings=args.stop_strings,
+    )
+
+    eos_id = tokenizer.eos_token_id
+
+    for step in range(args.max_new_tokens):
+        idx_cond = idx[:, -args.ctx_len:] if args.ctx_len > 0 else idx
+
+        logits = model_forward_logits(model, idx_cond)
+        logits = logits[:, -1, :][0]
+
+        if step < args.min_new_tokens:
+            if eos_id is not None and 0 <= eos_id < logits.numel():
+                logits[eos_id] = -float("inf")
+
+            for seq in stop_sequences:
+                if len(seq) == 1:
+                    tid = seq[0]
+                    if 0 <= tid < logits.numel():
+                        logits[tid] = -float("inf")
+
+        next_id = sample_next_token(
+            logits=logits,
+            generated_ids=generated_after_prompt,
+            tokenizer=tokenizer,
+            args=args,
+        )
+
+        next_tensor = torch.tensor(
+            [[next_id]],
+            dtype=torch.long,
+            device=idx.device,
+        )
+
+        idx = torch.cat([idx, next_tensor], dim=1)
+        generated_after_prompt.append(next_id)
+
+        full_ids = idx[0].tolist()
+
+        if step >= args.min_new_tokens:
+            if eos_id is not None and next_id == eos_id:
+                break
+
+            should_stop = False
+
+            for seq in stop_sequences:
+                if endswith_sequence(full_ids, seq):
+                    should_stop = True
+                    break
+
+            if should_stop:
+                break
+
+    return idx
+
+
+# ============================================================
+# CLI
+# ============================================================
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Inference script for Arithmetic-SLM using [IM_START]/[IM_END]."
+    )
+
+    p.add_argument(
+        "--model",
+        default="PhysiQuanty/Arithmetic-SLM",
+        help="HF model id or local path.",
+    )
+
+    p.add_argument(
+        "--prompt",
+        default="59 + 45 =",
+        help="Raw arithmetic prompt. With --no-think, inserted in chat template.",
+    )
+
+    p.add_argument(
+        "--no-think",
+        action="store_true",
+        help="Use production no-think template with [IM_START]/[IM_END].",
+    )
+
+    p.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+    )
+
+    p.add_argument(
+        "--dtype",
+        default="bfloat16",
+        choices=["bfloat16", "bf16", "float16", "fp16", "float32", "fp32"],
+    )
+
+    p.add_argument("--ctx-len", type=int, default=2048)
+    p.add_argument("--max-new-tokens", type=int, default=64)
+    p.add_argument("--min-new-tokens", type=int, default=1)
+
+    p.add_argument("--temperature", type=float, default=0.7)
+    p.add_argument("--top-k", type=int, default=40)
+    p.add_argument("--top-p", type=float, default=0.90)
+    p.add_argument("--min-p", type=float, default=0.0)
+    p.add_argument("--typical-p", type=float, default=1.0)
+
+    p.add_argument("--repetition-penalty", type=float, default=1.05)
+    p.add_argument("--frequency-penalty", type=float, default=0.10)
+    p.add_argument("--presence-penalty", type=float, default=0.0)
+    p.add_argument("--no-repeat-ngram-size", type=int, default=4)
+
+    p.add_argument("--seed", type=int, default=-1)
+
+    p.add_argument(
+        "--stop-string",
+        action="append",
+        default=None,
+        help="Additional stop string. Can be passed multiple times.",
+    )
+
+    p.add_argument(
+        "--bad-word",
+        action="append",
+        default=None,
+        help="Single-token word/token to ban. Can be passed multiple times.",
+    )
+
+    p.add_argument(
+        "--print-full",
+        action="store_true",
+        help="Print full prompt + completion.",
+    )
+
+    p.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        default=True,
+    )
+
+    args = p.parse_args()
+
+    if args.max_new_tokens <= 0:
+        raise ValueError("--max-new-tokens must be > 0")
+
+    if args.min_new_tokens < 0:
+        raise ValueError("--min-new-tokens must be >= 0")
+
+    if args.temperature < 0:
+        raise ValueError("--temperature must be >= 0")
+
+    if args.top_k < 0:
+        raise ValueError("--top-k must be >= 0")
+
+    if not (0 < args.top_p <= 1.0):
+        raise ValueError("--top-p must be in (0, 1]")
+
+    if args.min_p < 0:
+        raise ValueError("--min-p must be >= 0")
+
+    if not (0 < args.typical_p <= 1.0):
+        raise ValueError("--typical-p must be in (0, 1]")
+
+    if args.repetition_penalty <= 0:
+        raise ValueError("--repetition-penalty must be > 0")
+
+    if args.no_repeat_ngram_size < 0:
+        raise ValueError("--no-repeat-ngram-size must be >= 0")
+
+    stop_strings = [
+        IM_END,
+        IM_START,
+    ]
+
+    if args.stop_string:
+        stop_strings.extend(args.stop_string)
+
+    args.stop_strings = stop_strings
+
+    bad_words = []
+
+    if args.bad_word:
+        bad_words.extend(args.bad_word)
+
+    args.bad_words = bad_words
+
+    return args
+
+
+def main():
+    args = parse_args()
+    set_seed(args.seed)
+
+    dtype = get_dtype(args.dtype)
+
+    print(f"[INFO] model:       {args.model}")
+    print(f"[INFO] device:      {args.device}")
+    print(f"[INFO] dtype:       {args.dtype}")
+    print(f"[INFO] template:    {'no_think' if args.no_think else 'raw'}")
+    print(f"[INFO] IM_START:    {IM_START}")
+    print(f"[INFO] IM_END:      {IM_END}")
+    print(f"[INFO] NO_THINK:    {NO_THINK}")
+    print(f"[INFO] temperature: {args.temperature}")
+    print(f"[INFO] top_k:       {args.top_k}")
+    print(f"[INFO] top_p:       {args.top_p}")
+    print(f"[INFO] min_p:       {args.min_p}")
+    print(f"[INFO] typical_p:   {args.typical_p}")
+    print()
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model,
+        trust_remote_code=args.trust_remote_code,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model,
+        dtype=dtype,
+        trust_remote_code=args.trust_remote_code,
+    ).to(args.device)
+
+    model.eval()
+
+    prompt = build_prompt(args)
+
+    encoded = tokenizer(
+        prompt,
+        return_tensors="pt",
+        add_special_tokens=False,
+    )
+
+    encoded.pop("token_type_ids", None)
+
+    input_ids = encoded["input_ids"].to(args.device)
+
+    output_ids = generate_manual(
+        model=model,
+        tokenizer=tokenizer,
+        input_ids=input_ids,
+        args=args,
+    )
+
+    full_text = decode(tokenizer, output_ids[0].tolist())
+
+    if args.print_full:
+        print(full_text)
+        return
+
+    completion = extract_completion(full_text, prompt)
+    completion = strip_after_stop_text(completion, args.stop_strings)
+
+    print(completion)
+
+
+if __name__ == "__main__":
+    main()