Views
No views yet

1from transformers import LlavaForConditionalGeneration, AutoProcessor
2import torch
3from PIL import Image
4
5raw_model_name_or_path = "/保存的完整模型路径"
6model = LlavaForConditionalGeneration.from_pretrained(raw_model_name_or_path, device_map="cuda:0", torch_dtype=torch.bfloat16)
7processor = AutoProcessor.from_pretrained(raw_model_name_or_path)
8model.eval()
9
10def build_model_input(model, processor):
11 messages = [
12 {"role": "system", "content": "You are a helpful assistant."},
13 {"role": "user", "content": "<image>\n 你是一位有深度的网络图片解读者,擅长解读和描述网络图片。你能洞察图片中的细微之处,对图中的人物面部表情、文字信息、情绪流露和背景寓意具有超强的理解力,描述信息需要详细。"}
14 ]
15 prompt = processor.tokenizer.apply_chat_template(
16 messages, tokenize=False, add_generation_prompt=True
17 )
18 image = Image.open("01.PNG")
19 inputs = processor(text=prompt, images=image, return_tensors="pt", return_token_type_ids=False)
20
21 for tk in inputs.keys():
22 inputs[tk] = inputs[tk].to(model.device)
23 generate_ids = model.generate(**inputs, max_new_tokens=200)
24
25 generate_ids = [
26 oid[len(iids):] for oid, iids in zip(generate_ids, inputs.input_ids)
27 ]
28 gen_text = processor.batch_decode(generate_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False)[0]
29 return gen_text
30build_model_input(model, processor)