<reponame>Sunmingzhen/CogView
# -*- encoding: utf-8 -*-
'''
@File    :   datasets.py
@Time    :   2021/01/11 21:01:51
@Author  :   <NAME> 
@Contact :   <EMAIL>
'''

# here put the import lib
import os
import ipdb
import sys
import math
import random
from tqdm import tqdm
import logging


import numpy as np
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from torchvision import transforms as T
import pickle
from collections import namedtuple

import PIL
import lmdb
import json
from einops import rearrange

import sys
sys.path.append('..')
from utils import get_logger
from .unified_tokenizer import get_tokenizer
from .templates import TextCodeTemplate, TextFramesTemplate

logger = logging.getLogger(__name__)


class LMDBDataset(Dataset):
    def __init__(self, path, process_fn):
        self.env = lmdb.open(
            path,
            max_readers=32,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )
        self.process_fn = process_fn
        if not self.env:
            raise IOError('Cannot open lmdb dataset', path)

        with self.env.begin(write=False) as txn:
            self.length = int(txn.get('length'.encode('utf-8')).decode('utf-8'))

    def __len__(self):
        return self.length

    def __getitem__(self, idx):

        with self.env.begin(write=False) as txn:
            key = str(idx).encode('utf-8')

            row = pickle.loads(txn.get(key))

            return self.process_fn(row)

def get_dataset_by_type(dataset_type, path: str, args):
    Logger = get_logger()
    if dataset_type == "MsrvttTokensDataset":
        Logger.info(f"Loading MsrvttTokensDataset for train...")
        train_ds = MsrvttTokensDataset(path, 'trainval', args)
        valid_ds = MsrvttTokensDataset(path, 'test', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, valid_ds, None
    elif dataset_type == "WebvidTokensDataset":
        Logger.info(f"Loading WebvidTokensDataset for train...")
        train_ds = WebvidTokensDataset(path, 'train', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, None, None
    elif dataset_type == "WebvidFramesDataset":
        Logger.info(f"Loading WebvidFramesDataset for train...")
        train_ds = WebvidFramesDataset(path, 'train', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, None, None
    else:
        assert ValueError(f"No implemention for {dataset_type}")

def get_dataset_by_type_origin(dataset_type, path: str, args):
    tokenizer = get_tokenizer()
    if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
        ml = args.max_position_embeddings_finetune
    else:
        ml = args.max_position_embeddings

    def pad_to_len(ret):
        if len(ret) < ml:  # pad
            return np.concatenate((ret,
                                   np.array([tokenizer['[PAD]']] * (ml - len(ret)))),
                                  axis=0), len(ret)
        else:
            if len(ret) > ml:
                logger.warning('Out of max len, truncated.')
            return ret[:ml], ml

    def process_fn(row):
        ret, attention_mask_sep = pad_to_len(row.flatten())
        return {'text': ret,
                'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))}

    Logger = get_logger()
    if dataset_type == "MsrvttDataset":
        Logger.info(f"Loading MsrvttTokensDataset for train...")
        return MsrvttTokensDataset(path, 'train', args)
    elif dataset_type == "WebvidDataset":
        Logger.info(f"Loading WebvidTokensDataset for train...")
        return WebvidTokensDataset(path, 'train', args)
    else:
        assert ValueError(f"No implemention for {dataset_type}")

# def get_dataset_by_type(dataset_type, path: str, args, DS_CLASS=LMDBDataset):
#
#     tokenizer = get_tokenizer()
#     if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
#         ml = args.max_position_embeddings_finetune
#     else:
#         ml = args.max_position_embeddings
#
#     def pad_to_len(ret):
#
#         if len(ret) < ml: # pad
#             return np.concatenate((ret,
#                 np.array([tokenizer['[PAD]']] * (ml - len(ret)))),
#                 axis=0), len(ret)
#         else:
#             if len(ret) > ml:
#                 logger.warning('Out of max len, truncated.')
#             return ret[:ml], ml
#     Logger = get_logger()
#     if dataset_type == "MsrvttDataset":
#         Logger.info("Loading MsrvttTokensDataset...")
#         return MsrvttTokensDataset(path, args)
#     elif dataset_type == 'TokenizedDataset':
#         Logger.info("Loading TokenizedDataset...")
#         # already tokenized when saved
#         def process_fn(row):
#             ret, attention_mask_sep = pad_to_len(row.flatten())
#             return {'text': ret,
#                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))
#                 }
#
#     elif dataset_type == 'TextCodeDataset':
#         Logger.info("Loading TextCodeDataset...")
#         def process_fn(row):
#             text, code = row[0], row[1].flatten()
#             ret = TextCodeTemplate(text, code)
#             ret, attention_mask_sep = pad_to_len(ret)
#             return {'text': ret,
#                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))
#                 }
#     return DS_CLASS(path, process_fn)

class MsrvttTokensDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=128, text_len=35, truncate_captions=True,
                 shuffle=True):
        super(MsrvttTokensDataset).__init__()
        self.args = args
        self.shuffle = shuffle

        self.truncate_captions = truncate_captions

        self.text_len = text_len
        self.image_size = image_size

        # get video_ids
        self.frame_path = os.path.join(dataset_path, split)
        self.video_ids = os.listdir(self.frame_path)
        assert self.video_ids[0].endswith('.npy')
        self.is_tokens = True
        self.video_ids = [id[:-4] for id in self.video_ids]

        # get video_captions
        self.caption_path = os.path.join(dataset_path, 'caption.pkl')
        self.caption_dict = pickle.load(open(self.caption_path, 'rb'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': ret_tokens,
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        frame_tokens = torch.from_numpy(np.load(os.path.join(self.frame_path, '{}.npy'.format(video_id))))
        if frame_tokens.shape[0] != 2560:
            print("+" * 20 + "  " + video_id + "  " + "+" * 20)
        assert frame_tokens.shape[0] == 2560, "frame_tokens.shape[0] != 2560"

        # currently only predict 4 frames
        frame_num = 4
        frame_tokens = frame_tokens[:256 * frame_num]

        # get video captions by random choice
        captions = self.caption_dict[video_id]
        captions = list(filter(lambda t: len(t) > 0, captions))

        caption = np.random.choice(captions)

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)

class WebvidTokensDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=64, text_len=35, truncate_captions=True,
                 shuffle=True):
        super(WebvidTokensDataset).__init__()
        self.shuffle = shuffle

        self.truncate_captions = truncate_captions

        self.text_len = text_len
        self.image_size = image_size

        self.frame_path = os.path.join(dataset_path, 'webvid_tokens')
        self.json_path = dataset_path
        if split == 'train':
            self.frame_path = os.path.join(self.frame_path, 'train')
            self.json_path = os.path.join(self.json_path, 'webvid_2M_train.json')
            self.video_ids = os.listdir(self.frame_path)
        else:
            raise ValueError(f"No Implemention for {split}")

        # get video_ids
        assert self.video_ids[0].endswith('.npy')
        self.is_tokens = True
        self.video_ids = [id[:-4] for id in self.video_ids]

        # get video_captions
        self.json_dict = json.load(open(self.json_path, 'r'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            # logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': torch.from_numpy(np.int64(ret_tokens)),
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        self.logger.info(f'skip sample idx: {ind}!!')
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        frame_tokens = np.load(os.path.join(self.frame_path, '{}.npy'.format(video_id)))
        if frame_tokens.shape[0] != 2560:
            print("+" * 20 + "  " + video_id + "  " + "+" * 20)
        assert frame_tokens.shape[0] == 2560, "frame_tokens.shape[0] != 2560"

        # currently only predict 4 frames
        frame_num = 4
        frame_tokens = frame_tokens[:256 * frame_num]

        # get video captions by random choice
        caption = self.json_dict[video_id]['caption']

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)


class WebvidFramesDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=128, text_len=35, truncate_captions=True,
                 shuffle=True):
        super().__init__()
        self.shuffle = shuffle

        self.text_len = text_len
        self.image_size = image_size
        self.json_path = dataset_path
        self.truncate_captions = truncate_captions
        self.frame_path = os.path.join(dataset_path, 'webvid_frames_10')

        if split == 'train':
            self.frame_path = os.path.join(self.frame_path, 'train')
            self.json_path = os.path.join(self.json_path, 'webvid_2M_train.json')
            # get video_ids
            video_ids = os.listdir(self.frame_path)
            self.video_ids = video_ids * 20
            assert self.video_ids[0].endswith('.npy') is False
        else:
            raise ValueError(f"No Implemention for {split}")

        # get video_captions
        self.json_dict = json.load(open(self.json_path, 'r'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

        self.transform = transforms.Compose([
                                transforms.Resize((args.image_size, args.image_size),
                                                  interpolation=PIL.Image.BILINEAR),
                                transforms.ToTensor()
                            ])


    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            # logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': torch.from_numpy(np.int64(ret_tokens)),
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        self.logger.info(f'skip sample idx: {ind}!!')
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        imgs = []
        cur_path = os.path.join(self.frame_path, video_id)
        for i in range(10):
            img_path = os.path.join(cur_path, f"{video_id}_{i}.jpg")
            img = torchvision.datasets.folder.pil_loader(img_path)
            cur_img = self.transform(img).unsqueeze(0)
            imgs.append(cur_img)

        # predict 10 frames
        frame_num = 10
        frame_tokens = self.tokenizer.img_tokenizer.get_codebook_indices(torch.cat(imgs, dim=0)).flatten() # tokens: 10 * (8*8) -> frames, tokens_flatten
        # self.logger.info(f"frame_tokens' len: {frame_tokens.shape}")
        # get video captions by random choice
        caption = self.json_dict[video_id]['caption']

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)

# -*- encoding: utf-8 -*-
'''
@File    :   datasets.py
@Time    :   2021/01/11 21:01:51
@Author  :   <NAME> 
@Contact :   <EMAIL>
'''

# here put the import lib
import os
import ipdb
import sys
import math
import random
from tqdm import tqdm
import logging


import numpy as np
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision import datasets, transforms
from torchvision import transforms as T
import pickle
from collections import namedtuple

import PIL
import lmdb
import json
from einops import rearrange

import sys
sys.path.append('..')
from utils import get_logger
from .unified_tokenizer import get_tokenizer
from .templates import TextCodeTemplate, TextFramesTemplate

logger = logging.getLogger(__name__)


class LMDBDataset(Dataset):
    def __init__(self, path, process_fn):
        self.env = lmdb.open(
            path,
            max_readers=32,
            readonly=True,
            lock=False,
            readahead=False,
            meminit=False,
        )
        self.process_fn = process_fn
        if not self.env:
            raise IOError('Cannot open lmdb dataset', path)

        with self.env.begin(write=False) as txn:
            self.length = int(txn.get('length'.encode('utf-8')).decode('utf-8'))

    def __len__(self):
        return self.length

    def __getitem__(self, idx):

        with self.env.begin(write=False) as txn:
            key = str(idx).encode('utf-8')

            row = pickle.loads(txn.get(key))

            return self.process_fn(row)

def get_dataset_by_type(dataset_type, path: str, args):
    Logger = get_logger()
    if dataset_type == "MsrvttTokensDataset":
        Logger.info(f"Loading MsrvttTokensDataset for train...")
        train_ds = MsrvttTokensDataset(path, 'trainval', args)
        valid_ds = MsrvttTokensDataset(path, 'test', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, valid_ds, None
    elif dataset_type == "WebvidTokensDataset":
        Logger.info(f"Loading WebvidTokensDataset for train...")
        train_ds = WebvidTokensDataset(path, 'train', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, None, None
    elif dataset_type == "WebvidFramesDataset":
        Logger.info(f"Loading WebvidFramesDataset for train...")
        train_ds = WebvidFramesDataset(path, 'train', args)
        Logger.info(f"t=Train data {len(train_ds)}.")
        return train_ds, None, None
    else:
        assert ValueError(f"No implemention for {dataset_type}")

def get_dataset_by_type_origin(dataset_type, path: str, args):
    tokenizer = get_tokenizer()
    if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
        ml = args.max_position_embeddings_finetune
    else:
        ml = args.max_position_embeddings

    def pad_to_len(ret):
        if len(ret) < ml:  # pad
            return np.concatenate((ret,
                                   np.array([tokenizer['[PAD]']] * (ml - len(ret)))),
                                  axis=0), len(ret)
        else:
            if len(ret) > ml:
                logger.warning('Out of max len, truncated.')
            return ret[:ml], ml

    def process_fn(row):
        ret, attention_mask_sep = pad_to_len(row.flatten())
        return {'text': ret,
                'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))}

    Logger = get_logger()
    if dataset_type == "MsrvttDataset":
        Logger.info(f"Loading MsrvttTokensDataset for train...")
        return MsrvttTokensDataset(path, 'train', args)
    elif dataset_type == "WebvidDataset":
        Logger.info(f"Loading WebvidTokensDataset for train...")
        return WebvidTokensDataset(path, 'train', args)
    else:
        assert ValueError(f"No implemention for {dataset_type}")

# def get_dataset_by_type(dataset_type, path: str, args, DS_CLASS=LMDBDataset):
#
#     tokenizer = get_tokenizer()
#     if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
#         ml = args.max_position_embeddings_finetune
#     else:
#         ml = args.max_position_embeddings
#
#     def pad_to_len(ret):
#
#         if len(ret) < ml: # pad
#             return np.concatenate((ret,
#                 np.array([tokenizer['[PAD]']] * (ml - len(ret)))),
#                 axis=0), len(ret)
#         else:
#             if len(ret) > ml:
#                 logger.warning('Out of max len, truncated.')
#             return ret[:ml], ml
#     Logger = get_logger()
#     if dataset_type == "MsrvttDataset":
#         Logger.info("Loading MsrvttTokensDataset...")
#         return MsrvttTokensDataset(path, args)
#     elif dataset_type == 'TokenizedDataset':
#         Logger.info("Loading TokenizedDataset...")
#         # already tokenized when saved
#         def process_fn(row):
#             ret, attention_mask_sep = pad_to_len(row.flatten())
#             return {'text': ret,
#                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))
#                 }
#
#     elif dataset_type == 'TextCodeDataset':
#         Logger.info("Loading TextCodeDataset...")
#         def process_fn(row):
#             text, code = row[0], row[1].flatten()
#             ret = TextCodeTemplate(text, code)
#             ret, attention_mask_sep = pad_to_len(ret)
#             return {'text': ret,
#                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep))
#                 }
#     return DS_CLASS(path, process_fn)

class MsrvttTokensDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=128, text_len=35, truncate_captions=True,
                 shuffle=True):
        super(MsrvttTokensDataset).__init__()
        self.args = args
        self.shuffle = shuffle

        self.truncate_captions = truncate_captions

        self.text_len = text_len
        self.image_size = image_size

        # get video_ids
        self.frame_path = os.path.join(dataset_path, split)
        self.video_ids = os.listdir(self.frame_path)
        assert self.video_ids[0].endswith('.npy')
        self.is_tokens = True
        self.video_ids = [id[:-4] for id in self.video_ids]

        # get video_captions
        self.caption_path = os.path.join(dataset_path, 'caption.pkl')
        self.caption_dict = pickle.load(open(self.caption_path, 'rb'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': ret_tokens,
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        frame_tokens = torch.from_numpy(np.load(os.path.join(self.frame_path, '{}.npy'.format(video_id))))
        if frame_tokens.shape[0] != 2560:
            print("+" * 20 + "  " + video_id + "  " + "+" * 20)
        assert frame_tokens.shape[0] == 2560, "frame_tokens.shape[0] != 2560"

        # currently only predict 4 frames
        frame_num = 4
        frame_tokens = frame_tokens[:256 * frame_num]

        # get video captions by random choice
        captions = self.caption_dict[video_id]
        captions = list(filter(lambda t: len(t) > 0, captions))

        caption = np.random.choice(captions)

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)

class WebvidTokensDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=64, text_len=35, truncate_captions=True,
                 shuffle=True):
        super(WebvidTokensDataset).__init__()
        self.shuffle = shuffle

        self.truncate_captions = truncate_captions

        self.text_len = text_len
        self.image_size = image_size

        self.frame_path = os.path.join(dataset_path, 'webvid_tokens')
        self.json_path = dataset_path
        if split == 'train':
            self.frame_path = os.path.join(self.frame_path, 'train')
            self.json_path = os.path.join(self.json_path, 'webvid_2M_train.json')
            self.video_ids = os.listdir(self.frame_path)
        else:
            raise ValueError(f"No Implemention for {split}")

        # get video_ids
        assert self.video_ids[0].endswith('.npy')
        self.is_tokens = True
        self.video_ids = [id[:-4] for id in self.video_ids]

        # get video_captions
        self.json_dict = json.load(open(self.json_path, 'r'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            # logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': torch.from_numpy(np.int64(ret_tokens)),
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        self.logger.info(f'skip sample idx: {ind}!!')
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        frame_tokens = np.load(os.path.join(self.frame_path, '{}.npy'.format(video_id)))
        if frame_tokens.shape[0] != 2560:
            print("+" * 20 + "  " + video_id + "  " + "+" * 20)
        assert frame_tokens.shape[0] == 2560, "frame_tokens.shape[0] != 2560"

        # currently only predict 4 frames
        frame_num = 4
        frame_tokens = frame_tokens[:256 * frame_num]

        # get video captions by random choice
        caption = self.json_dict[video_id]['caption']

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)


class WebvidFramesDataset(Dataset):
    def __init__(self, dataset_path, split, args,  image_size=128, text_len=35, truncate_captions=True,
                 shuffle=True):
        super().__init__()
        self.shuffle = shuffle

        self.text_len = text_len
        self.image_size = image_size
        self.json_path = dataset_path
        self.truncate_captions = truncate_captions
        self.frame_path = os.path.join(dataset_path, 'webvid_frames_10')

        if split == 'train':
            self.frame_path = os.path.join(self.frame_path, 'train')
            self.json_path = os.path.join(self.json_path, 'webvid_2M_train.json')
            # get video_ids
            video_ids = os.listdir(self.frame_path)
            self.video_ids = video_ids * 20
            assert self.video_ids[0].endswith('.npy') is False
        else:
            raise ValueError(f"No Implemention for {split}")

        # get video_captions
        self.json_dict = json.load(open(self.json_path, 'r'))

        self.logger = get_logger()
        self.tokenizer = get_tokenizer()
        if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings:
            self.ml = args.max_position_embeddings_finetune
        else:
            self.ml = args.max_position_embeddings

        self.transform = transforms.Compose([
                                transforms.Resize((args.image_size, args.image_size),
                                                  interpolation=PIL.Image.BILINEAR),
                                transforms.ToTensor()
                            ])


    def process_fn(self, row):
        ret = row.flatten()
        if len(ret) <= self.ml:  # pad
            ret_tokens = np.concatenate((ret,
                                   np.array([self.tokenizer['[PAD]']] * (self.ml - len(ret)))),
                                  axis=0)
            attention_mask_sep= len(ret)
        else:
            # logger.warning(f'length of {len(ret)} out of max len, truncated.')
            ret_tokens = ret[:self.ml]
            attention_mask_sep= self.ml

        return {'text': torch.from_numpy(np.int64(ret_tokens)),
                'loss_mask': np.array([1] * attention_mask_sep + [0] * (self.ml - attention_mask_sep))}

    def __len__(self):
        return len(self.video_ids)

    def random_sample(self):
        return self.__getitem__(np.random.randint(0, self.__len__() - 1))

    def sequential_sample(self, ind):
        if ind >= self.__len__() - 1:
            return self.__getitem__(0)
        return self.__getitem__(ind + 1)

    def skip_sample(self, ind):
        self.logger.info(f'skip sample idx: {ind}!!')
        if self.shuffle:
            return self.random_sample()
        return self.sequential_sample(ind=ind)

    def __getitem__(self, index):
        video_id = self.video_ids[index]

        # get video tokens of 10 frames
        imgs = []
        cur_path = os.path.join(self.frame_path, video_id)
        for i in range(10):
            img_path = os.path.join(cur_path, f"{video_id}_{i}.jpg")
            img = torchvision.datasets.folder.pil_loader(img_path)
            cur_img = self.transform(img).unsqueeze(0)
            imgs.append(cur_img)

        # predict 10 frames
        frame_num = 10
        frame_tokens = self.tokenizer.img_tokenizer.get_codebook_indices(torch.cat(imgs, dim=0)).flatten() # tokens: 10 * (8*8) -> frames, tokens_flatten
        # self.logger.info(f"frame_tokens' len: {frame_tokens.shape}")
        # get video captions by random choice
        caption = self.json_dict[video_id]['caption']

        # wrap txt&img tokens
        tokens = TextFramesTemplate(caption, frame_tokens, frame_num)
        return self.process_fn(tokens)
# -*- encoding: utf-8 -*- @File    :   datasets.py @Time    :   2021/01/11 21:01:51 @Author  :   <NAME>  @Contact :   <EMAIL> # here put the import lib # pad # def get_dataset_by_type(dataset_type, path: str, args, DS_CLASS=LMDBDataset): # #     tokenizer = get_tokenizer() #     if args.finetune and args.max_position_embeddings_finetune > args.max_position_embeddings: #         ml = args.max_position_embeddings_finetune #     else: #         ml = args.max_position_embeddings # #     def pad_to_len(ret): # #         if len(ret) < ml: # pad #             return np.concatenate((ret, #                 np.array([tokenizer['[PAD]']] * (ml - len(ret)))), #                 axis=0), len(ret) #         else: #             if len(ret) > ml: #                 logger.warning('Out of max len, truncated.') #             return ret[:ml], ml #     Logger = get_logger() #     if dataset_type == "MsrvttDataset": #         Logger.info("Loading MsrvttTokensDataset...") #         return MsrvttTokensDataset(path, args) #     elif dataset_type == 'TokenizedDataset': #         Logger.info("Loading TokenizedDataset...") #         # already tokenized when saved #         def process_fn(row): #             ret, attention_mask_sep = pad_to_len(row.flatten()) #             return {'text': ret, #                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep)) #                 } # #     elif dataset_type == 'TextCodeDataset': #         Logger.info("Loading TextCodeDataset...") #         def process_fn(row): #             text, code = row[0], row[1].flatten() #             ret = TextCodeTemplate(text, code) #             ret, attention_mask_sep = pad_to_len(ret) #             return {'text': ret, #                 'loss_mask':  np.array([1] * attention_mask_sep + [0] * (len(ret) - attention_mask_sep)) #                 } #     return DS_CLASS(path, process_fn) # get video_ids # get video_captions # pad # get video tokens of 10 frames # currently only predict 4 frames # get video captions by random choice # wrap txt&img tokens # get video_ids # get video_captions # pad # logger.warning(f'length of {len(ret)} out of max len, truncated.') # get video tokens of 10 frames # currently only predict 4 frames # get video captions by random choice # wrap txt&img tokens # get video_ids # get video_captions # pad # logger.warning(f'length of {len(ret)} out of max len, truncated.') # get video tokens of 10 frames # predict 10 frames # tokens: 10 * (8*8) -> frames, tokens_flatten # self.logger.info(f"frame_tokens' len: {frame_tokens.shape}") # get video captions by random choice # wrap txt&img tokens