Spaces:
Running
Running
import base64 | |
from io import BytesIO | |
import torch | |
from PIL import Image | |
from transformers import StoppingCriteria | |
from .constants import IMAGE_TOKEN_INDEX | |
def load_image_from_base64(image): | |
return Image.open(BytesIO(base64.b64decode(image))) | |
def process_images(images, image_processor, model_cfg): | |
return image_processor(images, return_tensors="pt")["pixel_values"] | |
def tokenizer_image_token( | |
prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None | |
): | |
prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")] | |
def insert_separator(X, sep): | |
return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1] | |
input_ids = [] | |
offset = 0 | |
if ( | |
len(prompt_chunks) > 0 | |
and len(prompt_chunks[0]) > 0 | |
and prompt_chunks[0][0] == tokenizer.bos_token_id | |
): | |
offset = 1 | |
input_ids.append(prompt_chunks[0][0]) | |
for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)): | |
input_ids.extend(x[offset:]) | |
if return_tensors is not None: | |
if return_tensors == "pt": | |
return torch.tensor(input_ids, dtype=torch.long) | |
raise ValueError(f"Unsupported tensor type: {return_tensors}") | |
return input_ids | |
def get_model_name_from_path(model_path): | |
model_path = model_path.strip("/") | |
model_paths = model_path.split("/") | |
if model_paths[-1].startswith("checkpoint-"): | |
return model_paths[-2] + "_" + model_paths[-1] | |
else: | |
return model_paths[-1] | |
class KeywordsStoppingCriteria(StoppingCriteria): | |
def __init__(self, keywords, tokenizer, input_ids): | |
self.keywords = keywords | |
self.keyword_ids = [] | |
for keyword in keywords: | |
cur_keyword_ids = tokenizer(keyword).input_ids | |
if ( | |
len(cur_keyword_ids) > 1 | |
and cur_keyword_ids[0] == tokenizer.bos_token_id | |
): | |
cur_keyword_ids = cur_keyword_ids[1:] | |
self.keyword_ids.append(torch.tensor(cur_keyword_ids)) | |
self.tokenizer = tokenizer | |
self.start_len = input_ids.shape[1] | |
def __call__( | |
self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs | |
) -> bool: | |
assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)" # TODO | |
offset = min(output_ids.shape[1] - self.start_len, 3) | |
self.keyword_ids = [ | |
keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids | |
] | |
for keyword_id in self.keyword_ids: | |
if output_ids[0, -keyword_id.shape[0] :] == keyword_id: | |
return True | |
outputs = self.tokenizer.batch_decode( | |
output_ids[:, -offset:], skip_special_tokens=True | |
)[0] | |
for keyword in self.keywords: | |
if keyword in outputs: | |
return True | |
return False | |