Blame for inference.py

"""Standalone download, base/adapter loading, and identical-input inference."""
import argparse
from contextlib import nullcontext
import hashlib
import importlib.metadata
import json
from pathlib import Path
import platform
import re
from unittest.mock import patch

MODELS = {
    "llama": ("Llama-3-Design11-LoRA", "meta-llama/Meta-Llama-3-8B-Instruct",
              "8afb486c1db24fe5011ec46dfbe5b5dccdb575c2"),
    "midm": ("Midm-Design11-LoRA", "K-intelligence/Midm-2.0-Mini-Instruct",
             "383eb221c52a32278f1985257b264ade8d982e60"),
}
TEST_QUERIES = [
    "온라인 쇼핑 고객이 비슷한 상품의 반복 추천으로 불편을 겪습니다. 핵심 문제와 UX 개선 대안 두 가지를 제안해 주세요.",
    "소규모 창업팀이 제한된 예산으로 서비스 UX를 개선하려 합니다. 고객 조사와 화면 개선의 우선순위를 어떻게 정할까요?",
    "패키지 디자인 프로젝트에서 승인 기준이 자주 바뀝니다. 재작업을 줄일 검토 순서와 변경 관리 방법을 제안해 주세요.",
]


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo-id", default="pibuchi/IDINNOLAB-LLMS")
    parser.add_argument("--revision", default="main", help="Use a commit SHA for reproducibility")
    parser.add_argument("--model", choices=MODELS, default="midm")
    parser.add_argument("--variant", choices=["base", "best", "both"], default="both")
    parser.add_argument("--query", help="Custom query; omitted = three synthetic execution-test queries")
    parser.add_argument("--max-new-tokens", type=int, default=4096)
    parser.add_argument("--device", choices=["cuda", "cpu"], default="cuda")
    parser.add_argument("--precision", choices=["nf4", "bf16", "fp32"], default="nf4")
    parser.add_argument("--cache-dir", type=Path, help="Empty separate cache for independent download checks")
    parser.add_argument("--output", type=Path, help="Local JSON execution record; never uploaded automatically")
    args = parser.parse_args()
    if args.max_new_tokens < 1 or (args.device == "cpu" and args.precision != "fp32"):
        parser.error("Use positive token limit; CPU requires --precision fp32")
    return args


def main():
    args = parse_args()
    import torch
    import safetensors
    import transformers.modeling_utils as modeling_utils
    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set_seed
    from huggingface_hub import HfApi, snapshot_download
    from peft import PeftConfig, PeftModel
    if args.device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA unavailable; install the CUDA PyTorch wheel or select CPU/fp32")
    if args.precision == "bf16" and args.device == "cuda" and not torch.cuda.is_bf16_supported():
        raise RuntimeError("BF16 unsupported; select nf4")
    set_seed(42)
    folder, base_id, base_revision = MODELS[args.model]
    resolved = HfApi().repo_info(args.repo_id, revision=args.revision).sha
    snapshot = Path(snapshot_download(args.repo_id, revision=resolved, cache_dir=args.cache_dir,
                                      allow_patterns=[f"{folder}/*"]))
    adapter = snapshot / folder
    config = PeftConfig.from_pretrained(adapter)
    if config.base_model_name_or_path != base_id or config.revision != base_revision:
        raise ValueError("Adapter base/revision differs from the released model")
    tokenizer = AutoTokenizer.from_pretrained(adapter)
    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token
    dtype = torch.float32 if args.precision == "fp32" else (
        torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16)
    kwargs = {"revision": base_revision, "cache_dir": args.cache_dir, "dtype": dtype,
              "device_map": {"": 0} if args.device == "cuda" else {"": "cpu"}, "attn_implementation": "sdpa"}
    if args.precision == "nf4":
        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True,
            bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=dtype)
    original = safetensors.safe_open
    def pread(filename, framework, device="cpu", **options):
        return original(filename, framework=framework, device=device, **{**options, "backend": "pread"})
    with patch.object(modeling_utils, "safe_open", pread) if platform.system() == "Windows" else nullcontext():
        base = AutoModelForCausalLM.from_pretrained(base_id, **kwargs)
    model = PeftModel.from_pretrained(base, adapter)
    model.eval()
    model.config.use_cache = True
    queries = [args.query] if args.query else TEST_QUERIES
    results = []
    for index, query in enumerate(queries):
        messages = [{"role": "system", "content": "디자인 전략을 검토하고 자연스러운 한국어로만 답하세요."},
                    {"role": "user", "content": query}]
        inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True,
            date_string="12 Oct 2026", return_dict=True, return_tensors="pt").to(args.device)
        ids = inputs["input_ids"][0].tolist()
        if len(ids) + args.max_new_tokens > model.config.max_position_embeddings:
            raise ValueError("Prompt plus output exceeds model context capacity")
        digest = hashlib.sha256(json.dumps(ids).encode("utf-8")).hexdigest()
        variants = ["base", "best"] if args.variant == "both" else [args.variant]
        for variant in variants:
            with model.disable_adapter() if variant == "base" else nullcontext():
                with torch.inference_mode():
                    generated = model.generate(**inputs, max_new_tokens=args.max_new_tokens,
                        do_sample=False, repetition_penalty=1.0, pad_token_id=tokenizer.pad_token_id)
            tokens = generated[0, len(ids):]
            answer = tokenizer.decode(tokens, skip_special_tokens=True).strip()
            results.append({"case": index + 1, "variant": variant, "query": query,
                "answer": answer, "input_ids_sha256": digest, "generated_tokens": len(tokens),
                "hit_token_limit": len(tokens) >= args.max_new_tokens,
                "nonempty": bool(answer), "contains_korean": bool(re.search("[가-힣]", answer)),
                "contains_null_character": "\x00" in answer})
    report = {"repo_id": args.repo_id, "adapter_revision": resolved, "base_model": base_id,
        "base_revision": base_revision, "python": platform.python_version(), "platform": platform.platform(),
        "versions": {name: importlib.metadata.version(name) for name in
                     ["torch", "transformers", "peft", "accelerate", "bitsandbytes", "huggingface-hub", "safetensors"]},
        "device": args.device, "precision": args.precision, "dtype": str(dtype),
        "date_string": "12 Oct 2026", "seed": 42, "max_new_tokens": args.max_new_tokens,
        "cache_dir": str(args.cache_dir) if args.cache_dir else None, "results": results,
        "scope": "Execution check with synthetic queries; not a design-quality benchmark"}
    content = json.dumps(report, ensure_ascii=False, indent=2)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(content, encoding="utf-8")
    print(content)


if __name__ == "__main__":
    main()