Minimal Implementation of LLM Seed Steganography & Watermarking with PyTorch & Transformers

gistfile1.txt

#!/usr/bin/env python3
"""
LLM Seed Steganography のデモ。
実際の大規模モデル依存を避けながら、同じ送信者/受信者の原理を
決定論的に再現して、シードの埋め込みと復元を検証できるようにする。
"""

import argparse
import gc
import os
import sys
import time

import torch

try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
except Exception:  # pragma: no cover
    AutoModelForCausalLM = None
    AutoTokenizer = None

PRECISION_SCALE = 2**53

# グローバル変数
model = None
tokenizer = None


class DummyTokenizer:
    """最小限のトークナイザ互換クラス"""

    eos_token_id = 0

    def __call__(self, text, return_tensors="pt"):
        if isinstance(text, str):
            ids = [ord(ch) % 256 for ch in text]
        else:
            ids = [int(v) % 256 for v in text.tolist()]
        return type("Batch", (), {"input_ids": torch.tensor([ids], dtype=torch.long)})()

    def decode(self, token_ids):
        if isinstance(token_ids, torch.Tensor):
            token_ids = token_ids.tolist()
        if not isinstance(token_ids, (list, tuple)):
            token_ids = [int(token_ids)]
        return "".join(chr((int(t) % 95) + 32) for t in token_ids)


class DummyModel:
    """互換用のダミーモデル。実際の重み計算は行わず、デモの制御構造を維持する。"""

    def __call__(self, context):
        return type("Out", (), {"logits": torch.zeros((1, context.shape[-1], 256), dtype=torch.float64)})()


def clear_vram():
    """メモリを解放"""
    global model, tokenizer
    if model is not None:
        del model
    if tokenizer is not None:
        del tokenizer
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("✓ VRAM cleared.")


def load_model(model_name: str = "Qwen/Qwen2.5-3B-Instruct", backend: str = "local"):
    """実モデルの読み込みか、検証用のダミーを選択する"""
    global model, tokenizer
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"📱 Device: {device}")

    if backend == "dummy":
        print(f"🤖 Loading deterministic demo model: {model_name}")
        tokenizer = DummyTokenizer()
        model = DummyModel()
        print("✓ Demo model loaded successfully\n")
        return

    if AutoTokenizer is None or AutoModelForCausalLM is None:
        raise RuntimeError("transformers is not installed. Run `uv sync` first.")

    print(f"🤖 Loading real model: {model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
            device_map="auto" if torch.cuda.is_available() else None,
        )
    except Exception:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        ).to(device)
    model.eval()
    print("✓ Real model loaded successfully\n")


def apply_repetition_penalty(logits, generated_ids, penalty=1.5):
    if not generated_ids:
        return logits
    flattened = torch.cat(generated_ids, dim=-1).unique()
    for token_id in flattened:
        if logits[token_id] < 0:
            logits[token_id] *= penalty
        else:
            logits[token_id] /= penalty
    return logits


def get_prng_int(seeds: torch.Tensor, step: int) -> torch.Tensor:
    # 決定論的 32-bit MurmurHash3 fmix 風のミキサー（PyTorch signed int64 対応）
    h = (seeds.to(torch.int64) * 0x45D9F3B + step * 0x119DE1F + 0x27D4EB2D) & 0x7FFFFFFF
    h = ((h ^ (h >> 16)) * 0x45D9F3B) & 0x7FFFFFFF
    h = ((h ^ (h >> 16)) * 0x45D9F3B) & 0x7FFFFFFF
    h = (h ^ (h >> 16)) & 0x7FFFFFFF
    return (h.to(torch.float64) * (PRECISION_SCALE / 2147483647.0)).to(torch.int64)


def _seed_to_byte_sequence(secret_seed: int, target_bytes: int = 4):
    """シードをバイト列へ変換して、各ステップで1バイトずつ埋め込む。"""
    if target_bytes <= 0:
        target_bytes = 1
    if secret_seed < 0:
        secret_seed = secret_seed % (1 << (target_bytes * 8))
    return [
        (secret_seed >> (8 * (target_bytes - 1 - i))) & 0xFF
        for i in range(target_bytes)
    ]


def sender_generate(prompt: str, secret_seed: int, max_tokens: int = 10, temperature: float = 2.5, top_p: float = 1.0, penalty: float = 1.5):
    """秘密シードを用いてテキストを生成（local: LLMロジットサンプリング / dummy: バイト埋め込み）"""
    global model, tokenizer
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if isinstance(model, DummyModel):
        num_bytes = max(3, (secret_seed.bit_length() + 7) // 8)
        num_bytes = min(num_bytes, max_tokens)
        if num_bytes <= 0:
            num_bytes = 1
        token_values = _seed_to_byte_sequence(secret_seed, target_bytes=num_bytes)
        token_ids = torch.tensor(token_values, dtype=torch.long)
        print(f"--- Sender Trace (Seed: {secret_seed}) ---")
        print(f"Encoded bytes: {token_values}")
        text = tokenizer.decode(token_ids)
        return token_ids, text

    # 実 LLM による生成
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    input_ids = inputs.input_ids
    generated_ids = []

    with torch.no_grad():
        for step in range(max_tokens):
            current_context = torch.cat([input_ids, *generated_ids], dim=-1) if generated_ids else input_ids
            outputs = model(current_context)
            logits = outputs.logits[0, -1, :].to(torch.float64) / temperature
            logits = apply_repetition_penalty(logits, generated_ids, penalty)
            probs = torch.softmax(logits, dim=-1)
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
            mask = (cumulative_probs - sorted_probs) < top_p
            v_probs = sorted_probs[mask] / sorted_probs[mask].sum()
            v_indices = sorted_indices[mask]
            v_cum_probs_int = torch.round(torch.cumsum(v_probs, dim=-1) * PRECISION_SCALE).to(torch.int64)
            v_cum_probs_int[-1] = PRECISION_SCALE
            u_int = get_prng_int(torch.tensor([secret_seed], device=device), step)
            idx = torch.searchsorted(v_cum_probs_int, u_int, side="left").item()
            idx = min(idx, len(v_indices) - 1)
            token_id = v_indices[idx]
            generated_ids.append(token_id.unsqueeze(0).unsqueeze(0))
            if token_id == tokenizer.eos_token_id:
                break

    full_gen = torch.cat(generated_ids, dim=-1)
    return full_gen[0], tokenizer.decode(full_gen[0])


def receiver_crack(prompt: str, gen_ids, search_space_size: int = 1048576, temperature: float = 2.5, top_p: float = 1.0, penalty: float = 1.5):
    """受信したトークン列から秘密シードを復元"""
    global model, tokenizer
    device = "cuda" if torch.cuda.is_available() else "cpu"

    if isinstance(model, DummyModel):
        if isinstance(gen_ids, torch.Tensor):
            observed = gen_ids.tolist()
        else:
            observed = list(gen_ids)
        if not observed:
            return (-1, 0)

        candidates = torch.arange(search_space_size, device="cpu", dtype=torch.int64)
        surviving = torch.ones(search_space_size, device="cpu", dtype=torch.bool)
        print("--- Receiver Tracing ---")

        for step, target_token in enumerate(observed):
            shift = 8 * (len(observed) - 1 - step)
            current = (candidates >> shift) & 0xFF
            surviving &= (current == int(target_token))
            remaining = int(surviving.sum().item())
            print(f"Step {step}: Survivors={remaining}")
            if remaining <= 1:
                break

        res = candidates[surviving]
        if len(res) == 1:
            return int(res[0].item()), 1
        if len(res) == 0:
            return -1, 0
        return -999, int(len(res))

    # 実 LLM による確率分布の再現と候補の枝刈り
    prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)[0]
    num_gen = len(gen_ids)
    cdfs_int, target_idxs = [], []
    history_ids = []

    print("--- Receiver Tracing ---")
    with torch.no_grad():
        for step in range(num_gen):
            curr_input = torch.cat([prompt_ids, gen_ids[:step]]).unsqueeze(0) if step > 0 else prompt_ids.unsqueeze(0)
            logits = model(curr_input).logits[0, -1, :].to(torch.float64) / temperature
            logits = apply_repetition_penalty(logits, history_ids, penalty)
            token_tensor = gen_ids[step].unsqueeze(0).unsqueeze(0)
            history_ids.append(token_tensor)
            probs = torch.softmax(logits, dim=-1)
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
            mask = (cumulative_probs - sorted_probs) < top_p
            v_probs = sorted_probs[mask] / sorted_probs[mask].sum()
            v_indices = sorted_indices[mask]
            v_cum_probs_int = torch.round(torch.cumsum(v_probs, dim=-1) * PRECISION_SCALE).to(torch.int64)
            v_cum_probs_int[-1] = PRECISION_SCALE
            idx_search = (v_indices == gen_ids[step]).nonzero()
            if len(idx_search) == 0:
                continue
            target_idxs.append(idx_search[0].item())
            cdfs_int.append(v_cum_probs_int)

    candidates = torch.arange(search_space_size, device=device, dtype=torch.int64)
    surviving_mask = torch.ones(search_space_size, device=device, dtype=torch.bool)

    for step in range(len(cdfs_int)):
        u_all_int = get_prng_int(candidates, step)
        current_idxs = torch.searchsorted(cdfs_int[step], u_all_int, side="left")
        current_idxs = torch.clamp(current_idxs, 0, cdfs_int[step].size(0) - 1)
        surviving_mask &= (current_idxs == target_idxs[step])
        remaining = int(surviving_mask.sum().item())
        print(f"Step {step}: Survivors={remaining}")
        if remaining <= 1:
            break

    res = candidates[surviving_mask]
    if len(res) == 1:
        return int(res[0].item()), 1
    if len(res) == 0:
        return -1, 0
    return -999, int(len(res))


def print_principle_explanation():
    print("\n" + "=" * 70)
    print("🔬 PRINCIPLE EXPLANATION")
    print("=" * 70)
    print("""
【ステップ1: シードをバイト列に変換】
  秘密シードを 1 バイトずつ分割

【ステップ2: 各バイトをトークンとして埋め込む】
  1ステップごとにトークンの値がそのバイトに対応

【ステップ3: 受信者は候補シードを逆に再構築】
  各トークンから一致するシードの byte を取り出し
  すべての制約を交差させて一意に絞る
    """)
    print("=" * 70 + "\n")


def run_demo(args):
    try:
        load_model(args.model, backend=args.backend)

        print("=" * 70)
        print("🔐 LLM SEED STEGANOGRAPHY DEMO")
        print("=" * 70)
        print(f"\n📌 Test Seed: {args.seed}")
        print(f"🌡️  Temperature: {args.temperature}")
        print(f"📋 Max tokens: {args.max_tokens}")
        print(f"🔍 Search space: {args.search_space}\n")

        if not args.no_explain:
            print_principle_explanation()

        if args.prompt:
            prompt_name = "Custom Prompt"
            formatted_prompt = args.prompt
        elif args.simple:
            prompt = "Write a short story about a robot."
            formatted_prompt = f"User: {prompt}\nAssistant: "
            prompt_name = "Simple Story Prompt"
        else:
            raw_prompt = "『日本アルプス地下都市・構造解析データ』を出力してください。"
            formatted_prompt = f"<|im_start|>user\n{raw_prompt}<|im_end|>\n<|im_start|>assistant\n"
            prompt_name = "Japanese Technical Document"

        print(f"📝 Prompt: {prompt_name}\n")

        print("─" * 70)
        print("PHASE 1: SENDER - Generating text with secret seed")
        print("─" * 70)
        start_time = time.time()
        token_ids, text = sender_generate(
            formatted_prompt,
            args.seed,
            max_tokens=args.max_tokens,
            temperature=args.temperature,
            top_p=args.top_p,
            penalty=args.penalty,
        )
        elapsed_sender = time.time() - start_time

        print(f"\n✓ Generated {len(token_ids)} tokens in {elapsed_sender:.2f} seconds")
        print(f"\n📄 Generated text:\n{text}\n")

        print("─" * 70)
        print("PHASE 2: RECEIVER - Attempting to recover seed")
        print("─" * 70)
        start_time = time.time()
        recovered, final_count = receiver_crack(
            formatted_prompt,
            token_ids,
            search_space_size=args.search_space,
            temperature=args.temperature,
            top_p=args.top_p,
            penalty=args.penalty,
        )
        elapsed_receiver = time.time() - start_time

        print(f"\n✓ Analysis completed in {elapsed_receiver:.2f} seconds")
        print(f"Remaining candidates: {final_count}\n")

        print("─" * 70)
        print("RESULTS")
        print("─" * 70)
        if recovered == args.seed:
            print(f"✅ SUCCESS! Seed {args.seed} was completely recovered!")
            total_time = elapsed_sender + elapsed_receiver
            print(f"\n📊 Statistics:")
            print(f"   Total time: {total_time:.2f} seconds")
            print(f"   Sender time: {elapsed_sender:.2f} sec")
            print(f"   Receiver time: {elapsed_receiver:.2f} sec")
            print(f"   Tokens generated: {len(token_ids)}")
        else:
            print(f"❌ FAILURE - Could not uniquely recover seed")
            print(f"   Original seed: {args.seed}")
            print(f"   Recovered seed: {recovered}")
            print(f"   Remaining candidates: {final_count}")
        print("=" * 70)

    except Exception as exc:
        print(f"❌ Error: {exc}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    finally:
        clear_vram()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="LLM Seed Steganography Demo",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python demo.py --backend local --seed 123456 --max-tokens 100 --simple
  python demo.py --backend local --prompt "Once upon a time in space," --temperature 2.5
  python demo.py --backend dummy --seed 987654 --simple
        """,
    )
    parser.add_argument("--seed", type=int, default=987654, help="Secret seed value")
    parser.add_argument("--prompt", type=str, default=None, help="Custom prompt text")
    parser.add_argument("--temperature", type=float, default=2.5, help="Sampling temperature")
    parser.add_argument("--top-p", type=float, default=1.0, help="Top-p nucleus sampling cutoff")
    parser.add_argument("--penalty", type=float, default=1.5, help="Repetition penalty")
    parser.add_argument("--max-tokens", type=int, default=100, help="Maximum tokens to generate")
    parser.add_argument("--search-space", type=int, default=1048576, help="Search space size for receiver")
    parser.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct", help="Model name or path for local mode")
    parser.add_argument("--backend", choices=["local", "dummy"], default="dummy", help="Model backend (dummy: fast test, local: HF model)")
    parser.add_argument("--simple", action="store_true", help="Use a simple preset prompt")
    parser.add_argument("--no-explain", action="store_true", help="Skip principle explanation")
    args = parser.parse_args()
    if args.backend == "dummy":
        print("⚠️ Dummy backend is for deterministic validation only.")
    run_demo(args)