File size: 1,306 Bytes
9465d36 | 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 | import torch
from PIL import Image
from transformers import AutoModel, AutoTokenizer
torch.manual_seed(100)
model = AutoModel.from_pretrained('./MiniCPM-V-4_5', trust_remote_code=True, # or openbmb/MiniCPM-o-2_6
attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('./MiniCPM-V-4_5', trust_remote_code=True) # or openbmb/MiniCPM-o-2_6
ckpt_path = "/user/qinyihua/checkpoints/model_tunnel/llava1.5_downsample_mlp_sft/sft/job_27496_ckpt_5000"
state_dict = torch.load(ckpt_path, map_location="cpu")
image = Image.open('./test_xx.jpg').convert('RGB')
enable_thinking=False # If `enable_thinking=True`, the thinking mode is enabled.
# First round chat
question = "What is the landform in the picture?"
msgs = [{'role': 'user', 'content': [image, question]}]
answer = model.chat(
msgs=msgs,
tokenizer=tokenizer,
enable_thinking=enable_thinking
)
print(answer)
# # Second round chat, pass history context of multi-turn conversation
# msgs.append({"role": "assistant", "content": [answer]})
# msgs.append({"role": "user", "content": ["What should I pay attention to when traveling here?"]})
# answer = model.chat(
# msgs=msgs,
# tokenizer=tokenizer
# )
# print(answer)
|