Views
No views yet
1import requests
2import torch
3import transformers
4from PIL import Image
5
6from transformers.generation.streamers import TextStreamer
7from llava.constants import DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX
8from llava.conversation import conv_templates, SeparatorStyle
9from llava.model.llava_gpt2 import LlavaGpt2ForCausalLM
10from llava.train.dataset import tokenizer_image_token
11
12
13if __name__ == "__main__":
14 model_path = 'Hibernates/Hibernates-JP-1.3b-Max'
15 device = "cuda" if torch.cuda.is_available() else "cpu"
16 torch_dtype = torch.bfloat16 if device=="cuda" else torch.float32
17
18 model = LlavaGpt2ForCausalLM.from_pretrained(
19 model_path,
20 low_cpu_mem_usage=True,
21 use_safetensors=True,
22 torch_dtype=torch_dtype,
23 device_map=device,
24 )
25 tokenizer = transformers.AutoTokenizer.from_pretrained(
26 model_path,
27 model_max_length=1532,
28 padding_side="right",
29 use_fast=False,
30 )
31 model.eval()
32
33 conv_mode = "v1"
34 conv = conv_templates[conv_mode].copy()
35
36 # image pre-process
37 image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
38 image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
39
40 if device == "cuda":
41 image_tensor = model.get_model().vision_tower.image_processor(image).unsqueeze(0).half().cuda().to(torch_dtype)
42 else:
43 image_tensor = model.get_model().vision_tower.image_processor(image).unsqueeze(0).to(torch_dtype)
44
45 # create prompt
46 # ユーザー: <image>\n{prompt}
47 prompt = "猫の隣には何がありますか？"
48 inp = DEFAULT_IMAGE_TOKEN + '\n' + prompt
49 conv.append_message(conv.roles[0], inp)
50 conv.append_message(conv.roles[1], None)
51 prompt = conv.get_prompt()
52
53 input_ids = tokenizer_image_token(
54 prompt,
55 tokenizer,
56 IMAGE_TOKEN_INDEX,
57 return_tensors='pt'
58 ).unsqueeze(0)
59 if device == "cuda":
60 input_ids = input_ids.to(device)
61
62 input_ids = input_ids[:, :-1] # </sep>がinputの最後に入るので削除する
63 stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
64 keywords = [stop_str]
65 streamer = TextStreamer(tokenizer, skip_prompt=True, timeout=20.0)
66
67 # predict
68 with torch.inference_mode():
69 output_id = model.generate(
70 inputs=input_ids,
71 images=image_tensor,
72 do_sample=False,
73 temperature=1.0,
74 top_p=1.0,
75 max_new_tokens=256,
76 streamer=streamer,
77 use_cache=True,
78 )
79 """猫の隣にはノートパソコンがあります。"""
80
81### 🔧 Advanced Usage
82
83<details>
84<summary>パラメータのカスタマイズ</summary>
85
86```python
87# 生成パラメータの調整
88model.generate(
89 temperature=0.8, # 創造性の制御
90 top_p=0.95, # サンプリングの多様性
91 max_new_tokens=1024 # 生成テキストの長さ
92)