#!/usr/bin/env python3
"""
Thanatos-27B — Hugging Face Transformers quickstart.

Loads the upstream Qwen 3.6 27B safetensors directly and runs a single
chat turn using its embedded chat template. Thanatos-27B is a
*wrapper* around that base, so for the transformers route there is nothing
to download from this repo — point at Qwen/Qwen3.6-27B and apply the same
system prompt the Modelfile uses.

Requirements:
    pip install --upgrade "transformers>=4.45" accelerate sentencepiece bitsandbytes

Memory:
    - bf16 full precision: ~54 GB VRAM (won't fit on a single 24 GB card).
    - 4-bit (bitsandbytes nf4): ~16 GB VRAM, runs on a 3090/4090 24 GB.
    - Fall back to device_map="auto" + bnb_4bit on consumer GPUs.

Usage:
    python transformers_quickstart.py
    python transformers_quickstart.py --no-4bit         # bf16, needs ~54 GB VRAM
    python transformers_quickstart.py --prompt "..."    # custom prompt
"""
from __future__ import annotations

import argparse
import sys

try:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
except ImportError as e:  # pragma: no cover
    sys.exit(
        f"Missing dependency: {e.name}. Install with:\n"
        "  pip install --upgrade 'transformers>=4.45' accelerate sentencepiece bitsandbytes"
    )


MODEL_ID = "Qwen/Qwen3.6-27B"

THANATOS_SYSTEM = (
    "You are Thanatos, a precise and capable assistant for reasoning, writing, "
    "coding, and long-form dialogue.\n\n"
    "Behavior rules:\n"
    "- Answer the user's actual request directly.\n"
    "- Be accurate, complete, and structured.\n"
    "- Think before answering, but do not get stuck in repetitive loops or "
    "meta-commentary.\n"
    "- If the request is ambiguous or incomplete, state what is missing and "
    "make the smallest reasonable assumption needed to continue.\n"
    "- If the user wants creative writing, preserve tone, continuity, and "
    "character consistency.\n"
    "- If the user wants analysis or technical help, prefer concrete steps, "
    "examples, and decisions over fluff.\n"
    "- Finish with a usable answer, not just planning."
)


def load(use_4bit: bool):
    kwargs: dict = {"device_map": "auto", "torch_dtype": torch.bfloat16}
    if use_4bit:
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
        kwargs.pop("torch_dtype", None)

    tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True, **kwargs)
    return tok, model


def generate(tok, model, prompt: str, max_new_tokens: int = 512) -> str:
    messages = [
        {"role": "system", "content": THANATOS_SYSTEM},
        {"role": "user", "content": prompt},
    ]
    inputs = tok.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)

    out = model.generate(
        inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=0.6,
        top_p=0.95,
        top_k=20,
        repetition_penalty=1.05,
    )
    return tok.decode(out[0][inputs.shape[-1]:], skip_special_tokens=True)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--prompt", default="Explain the Burrows-Wheeler transform in 200 words.")
    ap.add_argument(
        "--no-4bit",
        action="store_true",
        help="Disable 4-bit quantization (requires ~54 GB VRAM in bf16).",
    )
    ap.add_argument("--max-new-tokens", type=int, default=512)
    args = ap.parse_args()

    print(f"[load] {MODEL_ID} (4bit={'no' if args.no_4bit else 'yes'})")
    tok, model = load(use_4bit=not args.no_4bit)

    print(f"[gen]  prompt: {args.prompt!r}")
    print()
    print(generate(tok, model, args.prompt, args.max_new_tokens))


if __name__ == "__main__":
    main()