Upload inference.py with huggingface_hub

inference.py

@@ -0,0 +1,260 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import torch
+
+from qwen_vl_utils import process_vision_info
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    AutoModelForImageTextToText,
+)
+
+
+DEFAULT_SYSTEM_PROMPT = (
+    "你是一个图像鉴伪专家，擅长结合视觉，文字结合伪造特征分析手段鉴别输入图像的真假。"
+    "分析过程中，你会逐步分析，抽丝剥茧，找到图像伪造的蛛丝马迹，"
+    "最终给出专业的鉴别结果及分析。"
+)
+
+DEFAULT_USER_PROMPT = (
+    "请分析这张文档图片是否存在伪造或篡改风险，并输出一份专业、精炼、准确的防伪分析报告。"
+)
+
+
+def build_messages(image_path: str, system_prompt: str, user_prompt: str):
+    return [
+        {
+            "role": "system",
+            "content": [
+                {
+                    "type": "text",
+                    "text": system_prompt,
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image",
+                    "image": image_path,
+                },
+                {
+                    "type": "text",
+                    "text": user_prompt,
+                },
+            ],
+        },
+    ]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument(
+        "--model-name",
+        type=str,
+        default="vankey/DocShield-9B",
+    )
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Input image path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=DEFAULT_USER_PROMPT,
+    )
+    parser.add_argument(
+        "--system-prompt",
+        type=str,
+        default=DEFAULT_SYSTEM_PROMPT,
+    )
+
+    thinking_group = parser.add_mutually_exclusive_group()
+    thinking_group.add_argument(
+        "--thinking",
+        action="store_true",
+        help="Enable Qwen3.5 thinking mode.",
+    )
+    thinking_group.add_argument(
+        "--no-thinking",
+        action="store_true",
+        help="Disable Qwen3.5 thinking mode.",
+    )
+
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=1024,
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.0,
+    )
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=0.8,
+    )
+    parser.add_argument(
+        "--top-k",
+        type=int,
+        default=20,
+    )
+    parser.add_argument(
+        "--do-sample",
+        action="store_true",
+        help="Use sampling. If not set, greedy decoding is used.",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16", "fp32"],
+    )
+
+    return parser.parse_args()
+
+
+def get_torch_dtype(dtype: str):
+    if dtype == "bf16":
+        return torch.bfloat16
+    if dtype == "fp16":
+        return torch.float16
+    return torch.float32
+
+
+def main():
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image not found: {args.image}")
+
+    enable_thinking = False
+    if args.thinking:
+        enable_thinking = True
+    if args.no_thinking:
+        enable_thinking = False
+
+    torch_dtype = get_torch_dtype(args.dtype)
+
+    print("=" * 100)
+    print("Model:", args.model_name)
+    print("Image:", args.image)
+    print("Prompt:", args.prompt)
+    print("Enable thinking:", enable_thinking)
+    print("Max new tokens:", args.max_new_tokens)
+    print("dtype:", args.dtype)
+    print("=" * 100)
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name,
+        use_fast=True,
+        trust_remote_code=True,
+    )
+
+    processor = AutoProcessor.from_pretrained(
+        args.model_name,
+        trust_remote_code=True,
+    )
+
+    model = AutoModelForImageTextToText.from_pretrained(
+        args.model_name,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True,
+        device_map="auto",
+    )
+
+    model.eval()
+
+    messages = build_messages(
+        image_path=args.image,
+        system_prompt=args.system_prompt,
+        user_prompt=args.prompt,
+    )
+
+    text = processor.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=enable_thinking,
+    )
+
+    print("\n" + "=" * 100)
+    print("Rendered prompt preview:")
+    print(text[:2000])
+    print("=" * 100 + "\n")
+
+    image_inputs, video_inputs = process_vision_info(messages)
+
+    inputs = processor(
+        text=[text],
+        images=image_inputs,
+        videos=video_inputs,
+        padding=True,
+        return_tensors="pt",
+    )
+
+    inputs = inputs.to(model.device)
+
+    print("input_ids shape:", inputs["input_ids"].shape)
+    if "pixel_values" in inputs:
+        print("pixel_values shape:", inputs["pixel_values"].shape)
+    if "image_grid_thw" in inputs:
+        print("image_grid_thw:", inputs["image_grid_thw"])
+
+    generation_kwargs = {
+        "max_new_tokens": args.max_new_tokens,
+    }
+
+    if args.do_sample:
+        generation_kwargs.update(
+            {
+                "do_sample": True,
+                "temperature": args.temperature,
+                "top_p": args.top_p,
+                "top_k": args.top_k,
+            }
+        )
+    else:
+        generation_kwargs.update(
+            {
+                "do_sample": False,
+            }
+        )
+
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **inputs,
+            **generation_kwargs,
+        )
+
+    generated_ids_trimmed = [
+        out_ids[len(in_ids):]
+        for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)
+    ]
+
+    output_text = processor.batch_decode(
+        generated_ids_trimmed,
+        skip_special_tokens=False,
+        clean_up_tokenization_spaces=False,
+    )[0]
+
+    print("\n" + "=" * 100)
+    print("Model output:")
+    print(output_text)
+    print("=" * 100)
+
+
+if __name__ == "__main__":
+    main()