Views
No views yet

1from transformers import AutoProcessor
2from vllm import LLM, SamplingParams
3from qwen_vl_utils import process_vision_info
4
5# 模型路径
6model_path = "path/DianJin-OCR-R1/seal_sft"
7# 图片路径
8image_path = "example.jpg"
9
10# 工具1的结果
11tool1 = "IXMTD5JPXGG9FEG10N 发票专用章 湄潭县何彬私房菜店"
12# 工具2的结果
13tool2 = "上海鸿路何彬私房菜连锁店 发票专用章"
14
15instruction = "请识别图片中的印章抬头。"
16tools = """<tool>
17以下是其它工具对该印章的识别内容:
18{{
19 "ocr_tool_1": "{tool1}",
20 "ocr_tool_2": "{tool2}"
21}}
22</tool>
23"""
24
25llm = LLM(
26 model=model_path,
27 limit_mm_per_prompt={"image": 10, "video": 10},
28 gpu_memory_utilization=0.4,
29)
30processor = AutoProcessor.from_pretrained(model_path)
31
32messages = [
33 {
34 "role": "user",
35 "content": [
36 {"type": "image", "image": image_path},
37 {"type": "text", "text": instruction},
38 ],
39 },
40]
41
42prompt = processor.apply_chat_template(
43 messages,
44 tokenize=False,
45 add_generation_prompt=True,
46)
47image_inputs, video_inputs, _ = process_vision_info(messages, return_video_kwargs=True)
48
49mm_data = {}
50if image_inputs is not None:
51 mm_data["image"] = image_inputs
52
53sampling_params = SamplingParams(
54 temperature=0.0,
55 top_p=1.0,
56 repetition_penalty=1.05,
57 max_tokens=4096,
58 stop=["<tool>"],
59)
60
61llm_inputs = [
62 {
63 "prompt": prompt,
64 "multi_modal_data": mm_data
65 }
66]
67
68outputs = llm.generate(llm_inputs, sampling_params=sampling_params)
69think_content = outputs[0].outputs[0].text.strip()
70print("#" * 20 + " think " + "#" * 20)
71print(think_content)
72
73llm_inputs[0]["prompt"] = (
74 llm_inputs[0]["prompt"].strip()
75 + "\n"
76 + think_content
77 + "\n"
78 + tools.format(tool1=tool1, tool2=tool2)
79)
80sampling_params = SamplingParams(
81 temperature=0.0, top_p=1.0, repetition_penalty=1.05, max_tokens=4096
82)
83outputs = llm.generate(llm_inputs, sampling_params=sampling_params)
84rethink_content = outputs[0].outputs[0].text.strip()
85print("#" * 20 + " rethink " + "#" * 20)
86print(rethink_content)
87