simplecloud/VidChain-exercise

1from PIL import Image2from io import BytesIO3import base644import numpy as np5import torch6import decord7from transformers import StoppingCriteria8from vtimellm.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, IMAGE_SEGMENT_TOKEN_INDEX, DEFAULT_IMAGE_SEGMENT_TOKEN9 10 11def load_image_from_base64(image):12    return Image.open(BytesIO(base64.b64decode(image)))13 14 15def process_images(images, image_processor, model_cfg):16    return image_processor(images, return_tensors='pt')['pixel_values']17 18 19def tokenizer_image_token_bf(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):20 21    def insert_separator(X, sep):22        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]23 24    prompt_chunks_t = prompt.split(DEFAULT_IMAGE_TOKEN)25 26    if (len(prompt_chunks_t) > 1 and DEFAULT_IMAGE_SEGMENT_TOKEN in prompt_chunks_t[1]):27        # incase <video segment> is present 28        prompt_chunks_seg_t = prompt_chunks_t[1].split(DEFAULT_IMAGE_SEGMENT_TOKEN)29        prompt_t = [prompt_chunks_t[0]] + prompt_chunks_seg_t30 31        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt_t]32 33        input_ids = []34        offset = 035        # if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:36        #     offset = 137        #     input_ids.append(prompt_chunks[0][0])38 39        input_ids = input_ids + prompt_chunks[0] + ([image_token_index] * (offset + 1))40 41        offset = 142        # image segment token43        for x in insert_separator(prompt_chunks[1:], [IMAGE_SEGMENT_TOKEN_INDEX] * (offset + 1)): 44            input_ids.extend(x[offset:])45 46        if return_tensors is not None:47            if return_tensors == 'pt':48                return torch.tensor(input_ids, dtype=torch.long)49            raise ValueError(f'Unsupported tensor type: {return_tensors}')50 51        return input_ids52    53    elif (len(prompt_chunks_t) == 1 and DEFAULT_IMAGE_SEGMENT_TOKEN in prompt_chunks_t[0]):54        # Assumed no image token in such prompt55        prompt_chunks_seg_t = prompt_chunks_t[0].split(DEFAULT_IMAGE_SEGMENT_TOKEN)56        prompt_t = prompt_chunks_seg_t57 58        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt_t]59 60        input_ids = []61        offset = 062 63        # image segment token]64        offset = 165        input_ids.append(prompt_chunks[0][0])66        for x in insert_separator(prompt_chunks, [IMAGE_SEGMENT_TOKEN_INDEX] * (offset + 1)): input_ids.extend(x[offset:])67 68        if return_tensors is not None:69            if return_tensors == 'pt':70                return torch.tensor(input_ids, dtype=torch.long)71            raise ValueError(f'Unsupported tensor type: {return_tensors}')72 73        return input_ids74 75    else:76        prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split(DEFAULT_IMAGE_TOKEN)]77 78        input_ids = []79        offset = 080        if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:81            offset = 182            input_ids.append(prompt_chunks[0][0])83        elif tokenizer.name == "GLMTokenizer":84            offset = 285            input_ids = prompt_chunks[0][:2]86 87        for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):88            input_ids.extend(x[offset:])89 90        if return_tensors is not None:91            if return_tensors == 'pt':92                return torch.tensor(input_ids, dtype=torch.long)93            raise ValueError(f'Unsupported tensor type: {return_tensors}')94        95        return input_ids96 97 98def tokenizer_image_token(prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors=None):99    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split(DEFAULT_IMAGE_TOKEN)]100 101    def insert_separator(X, sep):102        return [ele for sublist in zip(X, [sep]*len(X)) for ele in sublist][:-1]103 104    input_ids = []105    offset = 0106    if len(prompt_chunks) > 0 and len(prompt_chunks[0]) > 0 and prompt_chunks[0][0] == tokenizer.bos_token_id:107        offset = 1108        input_ids.append(prompt_chunks[0][0])109    elif tokenizer.name == "GLMTokenizer":110        offset = 2111        input_ids = prompt_chunks[0][:2]112 113    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):114        input_ids.extend(x[offset:])115 116    if return_tensors is not None:117        if return_tensors == 'pt':118            return torch.tensor(input_ids, dtype=torch.long)119        raise ValueError(f'Unsupported tensor type: {return_tensors}')120    return input_ids121 122 123def get_model_name_from_path(model_path):124    model_path = model_path.strip("/")125    model_paths = model_path.split("/")126    if model_paths[-1].startswith('checkpoint-'):127        return model_paths[-2] + "_" + model_paths[-1]128    else:129        return model_paths[-1]130 131 132 133 134class KeywordsStoppingCriteria(StoppingCriteria):135    def __init__(self, keywords, tokenizer, input_ids):136        self.keywords = keywords137        self.keyword_ids = []138        for keyword in keywords:139            cur_keyword_ids = tokenizer(keyword).input_ids140            if len(cur_keyword_ids) > 1 and cur_keyword_ids[0] == tokenizer.bos_token_id:141                cur_keyword_ids = cur_keyword_ids[1:]142            self.keyword_ids.append(torch.tensor(cur_keyword_ids))143        self.tokenizer = tokenizer144        self.start_len = input_ids.shape[1]145 146    def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:147        assert output_ids.shape[0] == 1, "Only support batch size 1 (yet)"  # TODO148        offset = min(output_ids.shape[1] - self.start_len, 3)149        self.keyword_ids = [keyword_id.to(output_ids.device) for keyword_id in self.keyword_ids]150        for keyword_id in self.keyword_ids:151            if output_ids[0, -keyword_id.shape[0]:].equal(keyword_id):152                return True153        outputs = self.tokenizer.batch_decode(output_ids[:, -offset:], skip_special_tokens=True)[0]154        for keyword in self.keywords:155            if keyword in outputs:156                return True157        return False158 159def print_trainable_parameters(model):160    trainable_params = 0161    all_param = 0162    for _, param in model.named_parameters():163        all_param += param.numel()164        # print(_, param.requires_grad, param.numel())165        if param.requires_grad:166            trainable_params += param.numel()167    print(168        f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param:.2f}"169    )170 171class VideoExtractor():172    """Dataset for supervised fine-tuning."""173 174    def __init__(self, N=100):175        self.N = N176 177    def extract(self, data):178        video_path = data['video']179        id = data['id']180        181        try:182            video_reader = decord.VideoReader(video_path)183            total_frames = len(video_reader)184            start = 0185            end = total_frames - 1186 187            split = data.get('split', None)188            if split is not None:189                fps = video_reader.get_avg_fps()190                start = max(int(fps * split[0]), 0)191                end = min(int(fps * split[1]), total_frames - 1)192            sampled_indices = np.linspace(start, end, self.N, dtype=np.int32)193            sampled_frames = video_reader.get_batch(sampled_indices).asnumpy()194        except Exception as e:195            print(e)196            return None, torch.zeros(1)197        198        images = torch.from_numpy(sampled_frames.transpose((0, 3, 1, 2)))199        return id, images