vankey/DocShield-9B

1#!/usr/bin/env python32import argparse3import os4import torch5 6from qwen_vl_utils import process_vision_info7from transformers import (8    AutoProcessor,9    AutoTokenizer,10    AutoModelForImageTextToText,11)12 13 14DEFAULT_SYSTEM_PROMPT = (15    "你是一个图像鉴伪专家，擅长结合视觉，文字结合伪造特征分析手段鉴别输入图像的真假。"16    "分析过程中，你会逐步分析，抽丝剥茧，找到图像伪造的蛛丝马迹，"17    "最终给出专业的鉴别结果及分析。"18)19 20DEFAULT_USER_PROMPT = (21    "请分析这张文档图片是否存在伪造或篡改风险，并输出一份专业、精炼、准确的防伪分析报告。"22)23 24 25def build_messages(image_path: str, system_prompt: str, user_prompt: str):26    return [27        {28            "role": "system",29            "content": [30                {31                    "type": "text",32                    "text": system_prompt,33                }34            ],35        },36        {37            "role": "user",38            "content": [39                {40                    "type": "image",41                    "image": image_path,42                },43                {44                    "type": "text",45                    "text": user_prompt,46                },47            ],48        },49    ]50 51 52def parse_args():53    parser = argparse.ArgumentParser()54 55    parser.add_argument(56        "--model-name",57        type=str,58        default="vankey/DocShield-9B",59    )60    parser.add_argument(61        "--image",62        type=str,63        required=True,64        help="Input image path.",65    )66    parser.add_argument(67        "--prompt",68        type=str,69        default=DEFAULT_USER_PROMPT,70    )71    parser.add_argument(72        "--system-prompt",73        type=str,74        default=DEFAULT_SYSTEM_PROMPT,75    )76 77    thinking_group = parser.add_mutually_exclusive_group()78    thinking_group.add_argument(79        "--thinking",80        action="store_true",81        help="Enable Qwen3.5 thinking mode.",82    )83    thinking_group.add_argument(84        "--no-thinking",85        action="store_true",86        help="Disable Qwen3.5 thinking mode.",87    )88 89    parser.add_argument(90        "--max-new-tokens",91        type=int,92        default=1024,93    )94    parser.add_argument(95        "--temperature",96        type=float,97        default=0.0,98    )99    parser.add_argument(100        "--top-p",101        type=float,102        default=0.8,103    )104    parser.add_argument(105        "--top-k",106        type=int,107        default=20,108    )109    parser.add_argument(110        "--do-sample",111        action="store_true",112        help="Use sampling. If not set, greedy decoding is used.",113    )114    parser.add_argument(115        "--device",116        type=str,117        default="cuda",118    )119    parser.add_argument(120        "--dtype",121        type=str,122        default="bf16",123        choices=["bf16", "fp16", "fp32"],124    )125 126    return parser.parse_args()127 128 129def get_torch_dtype(dtype: str):130    if dtype == "bf16":131        return torch.bfloat16132    if dtype == "fp16":133        return torch.float16134    return torch.float32135 136 137def main():138    args = parse_args()139 140    if not os.path.exists(args.image):141        raise FileNotFoundError(f"Image not found: {args.image}")142 143    enable_thinking = False144    if args.thinking:145        enable_thinking = True146    if args.no_thinking:147        enable_thinking = False148 149    torch_dtype = get_torch_dtype(args.dtype)150 151    print("=" * 100)152    print("Model:", args.model_name)153    print("Image:", args.image)154    print("Prompt:", args.prompt)155    print("Enable thinking:", enable_thinking)156    print("Max new tokens:", args.max_new_tokens)157    print("dtype:", args.dtype)158    print("=" * 100)159 160    tokenizer = AutoTokenizer.from_pretrained(161        args.model_name,162        use_fast=True,163        trust_remote_code=True,164    )165 166    processor = AutoProcessor.from_pretrained(167        args.model_name,168        trust_remote_code=True,169    )170 171    model = AutoModelForImageTextToText.from_pretrained(172        args.model_name,173        torch_dtype=torch_dtype,174        trust_remote_code=True,175        device_map="auto",176    )177 178    model.eval()179 180    messages = build_messages(181        image_path=args.image,182        system_prompt=args.system_prompt,183        user_prompt=args.prompt,184    )185 186    text = processor.apply_chat_template(187        messages,188        tokenize=False,189        add_generation_prompt=True,190        enable_thinking=enable_thinking,191    )192 193    print("\n" + "=" * 100)194    print("Rendered prompt preview:")195    print(text[:2000])196    print("=" * 100 + "\n")197 198    image_inputs, video_inputs = process_vision_info(messages)199 200    inputs = processor(201        text=[text],202        images=image_inputs,203        videos=video_inputs,204        padding=True,205        return_tensors="pt",206    )207 208    inputs = inputs.to(model.device)209 210    print("input_ids shape:", inputs["input_ids"].shape)211    if "pixel_values" in inputs:212        print("pixel_values shape:", inputs["pixel_values"].shape)213    if "image_grid_thw" in inputs:214        print("image_grid_thw:", inputs["image_grid_thw"])215 216    generation_kwargs = {217        "max_new_tokens": args.max_new_tokens,218    }219 220    if args.do_sample:221        generation_kwargs.update(222            {223                "do_sample": True,224                "temperature": args.temperature,225                "top_p": args.top_p,226                "top_k": args.top_k,227            }228        )229    else:230        generation_kwargs.update(231            {232                "do_sample": False,233            }234        )235 236    with torch.no_grad():237        generated_ids = model.generate(238            **inputs,239            **generation_kwargs,240        )241 242    generated_ids_trimmed = [243        out_ids[len(in_ids):]244        for in_ids, out_ids in zip(inputs["input_ids"], generated_ids)245    ]246 247    output_text = processor.batch_decode(248        generated_ids_trimmed,249        skip_special_tokens=False,250        clean_up_tokenization_spaces=False,251    )[0]252 253    print("\n" + "=" * 100)254    print("Model output:")255    print(output_text)256    print("=" * 100)257 258 259if __name__ == "__main__":260    main()261