Views
No views yet
1
2from transformers import AutoModelForImageTextToText, AutoProcessor
3
4MODEL_ID = "..." # path
5# Load model and tokenizer
6model = AutoModelForImageTextToText.from_pretrained(
7 MODEL_ID,
8 torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
9 device_map="auto",
10 low_cpu_mem_usage=True,
11)
12processor = AutoProcessor.from_pretrained(MODEL_ID)
13
14# Format question example
15SYSTEM_PROMPT = "..."
16img = None
17
18conversation = [
19 {"role": "system", "content": [{"type": "text", "text": SYSTEM_PROMPT}]},
20 {
21 "role": "user",
22 "content": [
23 {"type": "image", "image": img},
24 {"type": "text", "text": "..."},
25 ],
26 },
27]
28
29# Generate output
30prompt_text = processor.apply_chat_template(
31 conversation, add_generation_prompt=True, tokenize=False
32)
33inputs = processor(text=prompt_text, images=img, return_tensors="pt")
34with torch.inference_mode():
35 gen_out = model.generate(
36 **inputs,
37 max_new_tokens=256,
38 do_sample=False,
39 return_dict_in_generate=True,
40 output_scores=False,
41 )
42 sequences = gen_out.sequences
43
44input_len = inputs["input_ids"].shape[1]
45gen_ids = sequences[0, input_len:]
46resp_text = processor.tokenizer.decode(
47 gen_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
48).strip()