q-future/Co-Instruct

1from PIL import Image2from io import BytesIO3import base644 5import torch6from transformers import StoppingCriteria7from mplug_owl2.constants import IMAGE_TOKEN_INDEX,DEFAULT_IMAGE_TOKEN8from icecream import ic9 10 11def load_image_from_base64(image):12    return Image.open(BytesIO(base64.b64decode(image)))13 14 15def expand2square(pil_img, background_color):16    width, height = pil_img.size17    if width == height:18        return pil_img19    elif width > height:20        result = Image.new(pil_img.mode, (width, width), background_color)21        result.paste(pil_img, (0, (width - height) // 2))22        return result23    else:24        result = Image.new(pil_img.mode, (height, height), background_color)25        result.paste(pil_img, ((height - width) // 2, 0))26        return result27 28 29def process_images(images, image_processor, model_cfg):30    image_aspect_ratio = getattr(model_cfg, "image_aspect_ratio", None)31    new_images = []32    if image_aspect_ratio == 'pad':33        for image in images:34            image = expand2square(image, tuple(int(x*255) for x in image_processor.image_mean))35            image = image_processor.preprocess(image, return_tensors='pt')['pixel_values'][0]36            new_images.append(image)37    else:38        return image_processor(images, return_tensors='pt')['pixel_values']39    if all(x.shape == new_images[0].shape for x in new_images):40        new_images = torch.stack(new_images, dim=0)41    return new_images42 43 44def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):45    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split(DEFAULT_IMAGE_TOKEN)]46 47    def insert_separator(X, sep):48        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]49 50    input_ids = []51    offset = 052    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:53        offset = 154        input_ids.append(prompt_chunks[0][0])55 56    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):57        input_ids.extend(x[offset:])58 59    if return_tensors is not None:60        if return_tensors == 'pt':61            return torch.tensor(input_ids, dtype=torch.long)62        raise ValueError(f'Unsupported tensor type: {return_tensors}')63    return input_ids64 65 66def get_model_name_from_path(model_path):67    model_path = model_path.strip("/")68    model_paths = model_path.split("/")69    if model_paths[-1].startswith('checkpoint-'):70        return model_paths[-2] + "_" + model_paths[-1]71    else:72        return model_paths[-1]73 74 75 76 77class KeywordsStoppingCriteria(StoppingCriteria):78    def __init__(self, keywords, tokenizer, input_ids):79        self.keywords = keywords80        self.keyword_ids = []81        self.max_keyword_len = 082        for keyword in keywords:83            cur_keyword_ids = tokenizer(keyword).input_ids84            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:85                cur_keyword_ids = cur_keyword_ids[1:]86            if len(cur_keyword_ids) > self.max_keyword_len:87                self.max_keyword_len = len(cur_keyword_ids)88            self.keyword_ids.append(torch.tensor(cur_keyword_ids))89        self.tokenizer = tokenizer90        self.start_len = input_ids.shape[1]91 92    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:93        assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)"  # TODO94        offset = min(output_ids.shape[1] - self.start_len, self.max_keyword_len)95        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]96        for keyword_id in self.keyword_ids:97            if (output_ids[0, -keyword_id.shape[0]:] == keyword_id).all():98                return True99        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]100        for keyword in self.keywords:101            if keyword in outputs:102                return True103        return False