File size: 4,113 Bytes
b4942cf |
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 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
from dataclasses import field, dataclass
from typing import Optional, Union, List
import torch
from PIL import Image
from ovis.model.modeling_ovis import Ovis
from ovis.util.constants import IMAGE_TOKEN
@dataclass
class RunnerArguments:
model_path: str
max_new_tokens: int = field(default=512)
do_sample: bool = field(default=False)
top_p: Optional[float] = field(default=None)
top_k: Optional[int] = field(default=None)
temperature: Optional[float] = field(default=None)
max_partition: int = field(default=9)
class OvisRunner:
def __init__(self, args: RunnerArguments):
self.model_path = args.model_path
self.dtype = torch.bfloat16
self.device = torch.cuda.current_device()
self.dtype = torch.bfloat16
self.model = Ovis.from_pretrained(self.model_path, torch_dtype=self.dtype, multimodal_max_length=8192)
self.model = self.model.eval().to(device=self.device)
self.eos_token_id = self.model.generation_config.eos_token_id
self.text_tokenizer = self.model.get_text_tokenizer()
self.pad_token_id = self.text_tokenizer.pad_token_id
self.visual_tokenizer = self.model.get_visual_tokenizer()
self.conversation_formatter = self.model.get_conversation_formatter()
self.image_placeholder = IMAGE_TOKEN
self.max_partition = args.max_partition
self.gen_kwargs = dict(
max_new_tokens=args.max_new_tokens,
do_sample=args.do_sample,
top_p=args.top_p,
top_k=args.top_k,
temperature=args.temperature,
repetition_penalty=None,
eos_token_id=self.eos_token_id,
pad_token_id=self.pad_token_id,
use_cache=True
)
def preprocess(self, inputs: List[Union[Image.Image, str]]):
# for single image and single text inputs, ensure image ahead
if len(inputs) == 2 and isinstance(inputs[0], str) and isinstance(inputs[1], Image.Image):
inputs = reversed(inputs)
# build query
query = ''
images = []
for data in inputs:
if isinstance(data, Image.Image):
query += self.image_placeholder + '\n'
images.append(data)
elif isinstance(data, str):
query += data.replace(self.image_placeholder, '')
elif data is not None:
raise RuntimeError(f'Invalid input type, expected `PIL.Image.Image` or `str`, but got {type(data)}')
# format conversation
prompt, input_ids, pixel_values = self.model.preprocess_inputs(
query, images, max_partition=self.max_partition)
attention_mask = torch.ne(input_ids, self.text_tokenizer.pad_token_id)
input_ids = input_ids.unsqueeze(0).to(device=self.device)
attention_mask = attention_mask.unsqueeze(0).to(device=self.device)
if pixel_values is not None:
pixel_values = [pixel_values.to(device=self.device, dtype=self.dtype)]
else:
pixel_values = [None]
return prompt, input_ids, attention_mask, pixel_values
def run(self, inputs: List[Union[Image.Image, str]]):
prompt, input_ids, attention_mask, pixel_values = self.preprocess(inputs)
output_ids = self.model.generate(
input_ids,
pixel_values=pixel_values,
attention_mask=attention_mask,
**self.gen_kwargs
)
output = self.text_tokenizer.decode(output_ids[0], skip_special_tokens=True)
input_token_len = input_ids.shape[1]
output_token_len = output_ids.shape[1]
response = dict(
prompt=prompt,
output=output,
prompt_tokens=input_token_len,
total_tokens=input_token_len + output_token_len
)
return response
if __name__ == '__main__':
runner_args = RunnerArguments(model_path='<model_path>')
runner = OvisRunner(runner_args)
image = Image.open('<image_path>')
text = '<prompt>'
response = runner.run([image, text])
print(response['output'])
|