Views
No views yet


pip install torch transformers accelerate pillowCUDA_VISIBLE_DEVICES=0.1import torch
2import transformers
3from transformers import AutoModelForCausalLM, AutoTokenizer
4from PIL import Image
5import warnings
6
7# disable some warnings
8transformers.logging.set_verbosity_error()
9transformers.logging.disable_progress_bar()
10warnings.filterwarnings('ignore')
11
12# set device
13device = 'cuda' # or cpu
14torch.set_default_device(device)
15
16# create model
17model = AutoModelForCausalLM.from_pretrained(
18 'BAAI/Bunny-v1_0-3B',
19 torch_dtype=torch.float16, # float32 for cpu
20 device_map='auto',
21 trust_remote_code=True)
22tokenizer = AutoTokenizer.from_pretrained(
23 'BAAI/Bunny-v1_0-3B',
24 trust_remote_code=True)
25
26# text prompt
27prompt = 'Why is the image funny?'
28text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{prompt} ASSISTANT:"
29text_chunks = [tokenizer(chunk).input_ids for chunk in text.split('<image>')]
30input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1], dtype=torch.long).unsqueeze(0).to(device)
31
32# image, sample images can be found in images folder
33image = Image.open('example_2.png')
34image_tensor = model.process_images([image], model.config).to(dtype=model.dtype, device=device)
35
36# generate
37output_ids = model.generate(
38 input_ids,
39 images=image_tensor,
40 max_new_tokens=100,
41 use_cache=True)[0]
42
43print(tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip())