【备忘】CogVLM2使用笔记_cogvlm2-llama3-chinese-chat-19b-int4

1. 单卡4090用int4量化模型

1.1 下载仓库以及int4量化模型

cd <where_you_want>
git clone https://github.com/THUDM/CogVLM2.git
cd CogVLM2
mkdir -p models #存放模型
1
2
3
4

从modelscope下载int4模型

#模型下载，随便建一个py文件，
from modelscope import snapshot_download
model_dir = snapshot_download('ZhipuAI/cogvlm2-llama3-chinese-chat-19B-int4')
1
2
3

下载后正常应该存放在~/.cache/modelscope/hub/下，把它mv到 <where_you_want>/CogVLM2/models下。

1.2 demo使用。

推理的demo在CogVLM2/basic_demo下。

1.2.1 命令行对话，使用cli_demo.py

• 修改指向本地模型

ROOT_PATH = "<where_you_want>/CogVLM2/"
MODEL_NAME = "cogvlm2-llama3-chinese-chat-19B-int4"
MODEL_PATH = os.path.join(ROOT_PATH, "models", MODEL_NAME) # 假设已经按照1.1放置模型文件
1
2
3

• 因为模型已经是int4量化模型，所以python cli_demo.py时不需要--quant 4参数。

1.2.2 网页对话，使用chainlit run web_demo.py

同理，需要修改模型地址，参考 >1.2.1

• 模型返回未结束时，不要打断
• 每轮对话只能针对一张图片，重新输入一张图片需要刷新页面重启对话。
  
• 推理效果还可以
  
  

2. 纯OCR识别体验

2.1 体验

因工作中有OCR的需求，所以特意试了一下：

• 纯手写体：出乎意料的好。
• 纯印刷体：未能达到100%，问题基本出现在长数字串的吞字（如身份证号），涉及隐私无法贴出。
• 手写+印刷体：手写体基本被忽略而不显示。
• 推理时长。实验为A4纸，纯文字材料，数量35张，平均耗时～15s/页。
  识别结果
  

2.2 批量OCR的代码（基于Cli_demo.py改）

• 代码放在<project_root_dir>/basic_demo/下
• 输入图片放在<project_root_dir>/basic_demo/images/下
• 输出结果以 .md 格式放在<project_root_dir>/ocr_output/下，原ocr_output下所有.md都会压缩成一个zip包。

import torch
import argparse
from PIL import Image
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import os
import glob
import re
import time
from datetime import datetime
import zipfile

# 设置调用参数
# Argument parser
parser = argparse.ArgumentParser(description="CogVLM2 Batch OCR Demo")
parser.add_argument('--quant', type=int, choices=[4, 8], help='Enable 4-bit or 8-bit precision loading', default=0)
# 新增root_path参数，表示项目的路径，默认值为basic_demo的上级目录，即CogVLM2
parser.add_argument('--root_path', type=str, help='Path to the Project root directory', default=os.path.join(os.getcwd(), os.pardir))
# 新增--imgpath参数，表示待OCR的原图存放目录，默认值为当前目录下的images子目录
default_imgpath = os.path.join(os.getcwd(), 'images')
parser.add_argument('--imgpath', type=str, help='Path to the image directory', default=default_imgpath)
args = parser.parse_args()

# 定义项目路径、模型路径、设备类型和精度设置等
ROOT_PATH = args.root_path
MODEL_NAME = "cogvlm2-llama3-chinese-chat-19B-int4"
MODEL_PATH = os.path.join(ROOT_PATH, "models", MODEL_NAME)
# MODEL_PATH = "THUDM/cogvlm2-llama3-chat-19B"
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
TORCH_TYPE = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8 else torch.float16

# Dynamically adjust precision based on model path
if 'int4' in MODEL_PATH:
    args.quant = 4

def check_cuda():
# Check GPU memory and adjust accordingly
# Check GPU memory
    if torch.cuda.is_available() and torch.cuda.get_device_properties(0).total_memory < 48 * 1024 ** 3 and not args.quant:
        print("GPU memory is less than 48GB. Please use cli_demo_multi_gpus.py or pass `--quant 4` or `--quant 8`.")
        exit()
    return True

def get_imgfile_list(path: str, extensions: list[str] = ['.jpg', '.jpeg', '.png', '.bmp', '.gif']) -> list[str]:
    """
    获取指定路径下的所有图片文件列表
    :param path: 文件夹路径
    :param extensions: 图片文件后缀列表
    :return: 文件列表
    """
    # 验证路径有效性并防止路径遍历攻击
    clean_path = os.path.normpath(path)
    if not os.path.exists(clean_path) or not clean_path.startswith(os.path.abspath(path)):
        raise ValueError(f"Invalid path: {path}")

    # 使用os.walk遍历文件夹，但直接通过文件后缀名进行匹配，而不是正则表达式
    matching_files = []
    for dirpath, dirnames, filenames in os.walk(clean_path):
        for filename in filenames:
            base_name, file_extension = os.path.splitext(filename)
            if file_extension.lower() in [ext.lower() for ext in extensions]:
                matching_files.append(os.path.join(dirpath, filename))

    return matching_files

def get_imgdata(img_file: str):
    #检查img_file的有效性
    image_path = img_file
    try:
        image_data = Image.open(image_path).convert('RGB')
    except IOError:
        print(f"Error: Invalid image file: {image_path}")
        exit(IOError)

    return image_data

def gen_summary(summary: dict, LogPath: str):
    # summary是一个字典变量，key为文件名，value为推理时长
    # 统计模型处理时间, 计算模型推理的平均时间
    total_time = 0
    file_count = len(summary)
    for k, v in summary.items():
        total_time += v
    avg_time = total_time / file_count
    print(f"Total {file_count} files OCRed, model inference avarage time: {avg_time:.4f} seconds")
    # 将summary写入日志文件
    logfilefullpath = os.path.join(LogPath, f'{datetime.now().strftime("%Y-%m-%d %H:%M:%S")}.log')
    with open(LogPath, 'a') as f:
        f.write(f"Total {file_count} files OCRed, model inference avarage time: {avg_time:.4f} seconds\n")
        for k, v in summary.items():
            f.write(f"{k}: {v:.4f} seconds\n")

def main():
    if not check_cuda():
        exit(1)
    # Load the tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_PATH,
        trust_remote_code=True
    )
    print(f"Step1, tokenizer initiated.args.imgpath={args.imgpath}")

    # Load the model with appropriate precision and quantization settings
    # Load the model
    if args.quant == 4:
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_PATH,
            torch_dtype=TORCH_TYPE,
            trust_remote_code=True,
            quantization_config=BitsAndBytesConfig(load_in_4bit=True),
            low_cpu_mem_usage=True
        ).eval()
    elif args.quant == 8:
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_PATH,
            torch_dtype=TORCH_TYPE,
            trust_remote_code=True,
            quantization_config=BitsAndBytesConfig(load_in_8bit=True),
            low_cpu_mem_usage=True
        ).eval()
    else:
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_PATH,
            torch_dtype=TORCH_TYPE,
            trust_remote_code=True
        ).eval().to(DEVICE)
    print(f'Step2: model loaded with {args.quant}')

    # 创建一个OCR_OUTPUT_PATH变量, 使它等于ROOT_PATH下的ocr_output文件夹
    # 判断OCR_OUTPUT_PATH文件夹是否存在
    # 如果不存在则创建它
    # 如果存在则将这个目录内的全部.md文件压缩成一个zip，并使用当前时间作为文件名，然后清空所有.md文件
    OCR_OUTPUT_PATH = os.path.join(ROOT_PATH, "ocr_output")
    if not os.path.exists(OCR_OUTPUT_PATH):
        os.mkdir(OCR_OUTPUT_PATH)
    else:
        # 压缩OCR_OUTPUT_PATH文件夹内的所有.md文件到一个zip文件中
        # 使用当前时间作为文件名
        zip_file_name = os.path.join(OCR_OUTPUT_PATH, f"{datetime.now().strftime('%Y-%M-%d-%H:%M:%S')}.zip")
        # 使用zipfile模块压缩文件夹内所有md文件
        zip_file = zipfile.ZipFile(zip_file_name, 'w', zipfile.ZIP_DEFLATED)
        for file in glob.glob(os.path.join(OCR_OUTPUT_PATH, "*.md")):
            zip_file.write(file, os.path.basename(file))
        #清空这些旧的md文件
        for file in glob.glob(os.path.join(OCR_OUTPUT_PATH, "*.md")):
            os.remove(file)
        zip_file.close()
        print(f"Step3: {OCR_OUTPUT_PATH} is cleaned and {zip_file_name} is created")
    
    prompts = '请OCR图片全文，水印文字请单独列出，返回结果请使用Markdown格式。'
    gen_kwargs = {
        "max_new_tokens": 2048,
        "pad_token_id": 128002,
        "top_k": 1,
    }
    summary = {}  # 统计模型处理时间
    img_filelist = get_imgfile_list(args.imgpath)
    # print(f'prompt={prompts}\n{img_filelist}')
    # 遍历处理img_filelist中的文件。每个文件用get_imagedata打开并转换后，输入大模型，得到结果后记录为输入文件同名的word文档。
    for img_file in img_filelist:
        mark_time = time.time()
        image_data = get_imgdata(img_file)
        convert_img_time = time.time() - mark_time
        filename = img_file.split('.')[0]
        input_by_model = model.build_conversation_input_ids(
            tokenizer,
            query=prompts,
            images=[image_data],
            template_version='chat'
        )
        # Structure input for model generation and specify generation parameters
        inputs = {
            'input_ids': input_by_model['input_ids'].unsqueeze(0).to(DEVICE),
            'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(DEVICE),
            'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(DEVICE),
            'images': [[input_by_model['images'][0].to(DEVICE).to(TORCH_TYPE)]] if image_data is not None else None,
        }
        # 生成大模型结果，并将结果存为与img_file同名的md文档。
        # print(f"\nimg:{img_file} OCRing...")
        print(f"\nImage:{os.path.basename(img_file)} OCRing", end='...')
        filename += '.md'
        filename = os.path.join(OCR_OUTPUT_PATH, os.path.basename(filename))
        with torch.no_grad():
            outputs = model.generate(**inputs, **gen_kwargs)
            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
            with open(filename, 'w', encoding='utf-8') as f:
                f.write(response)
        generate_time = time.time() - mark_time
        #在summary中追加一条记录，key为os.path.basename(filename)，value为generate_time
        summary[os.path.basename(filename)] = generate_time
        print(f"Done! Infer time: {generate_time:.4f} seconds")
        # print(f"Summary: \n \t\t convert_img_time = {convert_img_time} 秒 \n \t\t generate_time = {generate_time} 秒")
    gen_summary(summary, OCR_OUTPUT_PATH)

if __name__=="__main__":
    main()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195