code: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  from io import BytesIO
  
  import logging
  import warnings
  import string
  
  import numpy as np
  import torch
  import base64
  from torchvision import transforms
  
  from PIL import Image, ImageFile
  
  from data import data_utils
  from data.ofa_dataset import OFADataset
  
  ImageFile.LOAD_TRUNCATED_IMAGES = True
  ImageFile.MAX_IMAGE_PIXELS = None
  Image.MAX_IMAGE_PIXELS = None
  
  logger = logging.getLogger(__name__)
  warnings.filterwarnings("ignore", "(Possibly )?corrupt EXIF data", UserWarning)
  
  IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
  IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
  
  
  def collate(samples, pad_idx, eos_idx):
      if len(samples) == 0:
          return {}
  
      def merge(key):
          return data_utils.collate_tokens(
              [s[key] for s in samples],
              pad_idx,
              eos_idx=eos_idx,
          )
  
      id = np.array([s["id"] for s in samples])
      src_tokens = merge("source")
      src_lengths = torch.LongTensor([s["source"].ne(pad_idx).long().sum() for s in samples])
  
      # patch_images = torch.stack([sample['patch_image'] for sample in samples], dim=0)
      # patch_masks = torch.cat([sample['patch_mask'] for sample in samples])
  
      prev_output_tokens = None
      target = None
      if samples[0].get("target", None) is not None:
          target = merge("target")
          tgt_lengths = torch.LongTensor([s["target"].ne(pad_idx).long().sum() for s in samples])
          ntokens = tgt_lengths.sum().item()
  
          if samples[0].get("prev_output_tokens", None) is not None:
              prev_output_tokens = merge("prev_output_tokens")
      else:
          ntokens = src_lengths.sum().item()
  
      batch = {
          "id": id,
          "nsentences": len(samples),
          "ntokens": ntokens,
          "net_input": {
              "src_tokens": src_tokens,
              "src_lengths": src_lengths,
              # "patch_images": patch_images,
              # "patch_masks": patch_masks,
              "prev_output_tokens": prev_output_tokens
          },
          "target": target,
      }
  
      return batch
  
  
  class Rec_nextitemDataset(OFADataset):
      def __init__(
          self,
          split,
          dataset,
          bpe,
          src_dict,
          tgt_dict=None,
          max_src_length=128,
          max_tgt_length=30,
          # patch_image_size=224,
          # imagenet_default_mean_and_std=False,
          scst=False
      ):
          super().__init__(split, dataset, bpe, src_dict, tgt_dict)
          self.max_src_length = max_src_length
          self.max_tgt_length = max_tgt_length
          # self.patch_image_size = patch_image_size
          self.scst = scst
  
          self.transtab = str.maketrans({key: None for key in string.punctuation})
  
          # if imagenet_default_mean_and_std:
          #     mean = IMAGENET_DEFAULT_MEAN
          #     std = IMAGENET_DEFAULT_STD
          # else:
          #     mean = [0.5, 0.5, 0.5]
          #     std = [0.5, 0.5, 0.5]
  
          # self.patch_resize_transform = transforms.Compose([
          #     lambda image: image.convert("RGB"),
          #     transforms.Resize((patch_image_size, patch_image_size), interpolation=Image.BICUBIC),
          #     transforms.ToTensor(),
          #     transforms.Normalize(mean=mean, std=std),
          # ])
          print("self.max_tgt_length", self.max_tgt_length)
  
      def __getitem__(self, index):
          uniq_id, user_behavior, target_item, rating = self.dataset[index]
  
          while target_item.translate(self.transtab).strip() == "":
              uniq_id, user_behavior, target_item, rating = self.dataset[index]
  
          if len(user_behavior) >= self.max_src_length - 20:
              user_behavior = user_behavior[:self.max_src_length - 20]
              if user_behavior[-1] != ",":
                  user_behavior = ','.join(user_behavior.split(",")[:-1])
              else:
                  user_behavior = user_behavior[:-1]
  
          if self.split == 'train' and not self.scst:
              target_item = target_item.translate(self.transtab).strip()
              target_item_token_list = target_item.strip().split(" ")
              tgt_explain = ' '.join(target_item_token_list[:self.max_tgt_length])
          else:
              target_item = ' '.join(target_item.strip().split(" ")[:self.max_tgt_length])
              target_item_list = [target_item.translate(self.transtab).strip() for explain in target_item.strip().split('&&')]
              tgt_explain = '&&'.join(target_item_list)
  
          print("user_behavior", user_behavior)
          src_text = "If you liked " + user_behavior + \
              ", you will also like "
  
          assert len(src_text.split(" ")) <= self.max_src_length
          src_item = self.encode_text(src_text)
          tgt_item = self.encode_text(" {}".format(tgt_explain))
  
          src_item = torch.cat([self.bos_item, src_item, self.eos_item])
          target_item = torch.cat([tgt_item, self.eos_item])
          prev_output_item = torch.cat([self.bos_item, tgt_item])
          example = {
              "id": uniq_id,
              "source": src_item,
              # "patch_image": patch_image,
              # "patch_mask": patch_mask,
              "target": target_item,
              "prev_output_tokens": prev_output_item
          }
          return example
  
      def collater(self, samples, pad_to_length=None):
          """Merge a list of samples to form a mini-batch.
          Args:
              samples (List[dict]): samples to collate
          Returns:
              dict: a mini-batch with the following keys:
          """
          return collate(samples, pad_idx=self.pad, eos_idx=self.eos)
  
  
      def __getitem__2(self, index):
          # uniq_id, user_behavior, explaination, fea, opt = self.dataset[index]
          # print("user_behavior", user_behavior)
          # print("fea", fea, "opt", opt)
          # print("user_behavior", user_behavior)
          tgt_explain = "asdasss ssa"
          while tgt_explain.strip() != "":
              uniq_id, user_behavior, explaination, fea, opt = self.dataset[index]
              while explaination.translate(self.transtab).strip() == "":
                  uniq_id, user_behavior, explaination, fea, opt = self.dataset[index]
              print("explaination begin", explaination)
              tmp_user_beha = user_behavior.split(" Right now, ")
              len_context = len(tmp_user_beha[1].split(" "))
              behavior_list = tmp_user_beha[0].split(" ")[0: self.max_src_length - 40 - len_context]
              behavior_ = " ".join(behavior_list)
              if behavior_[-1] == ",":
                  behavior_ = behavior_[:-1] + '.'
              if behavior_[-1] != ".":
                  behavior_ = ','.join(behavior_.split(",")[:-1]) + '.'
  
              user_behavior = " right now, ".join([behavior_, tmp_user_beha[1]])
  
              user_behavior += \
                  " the user cares about {} and the item is {}.".format(fea, opt)
  
  
              # image = Image.open(BytesIO(base64.urlsafe_b64decode(image)))
              # patch_image = self.patch_resize_transform(image)
              # patch_mask = torch.tensor([True])
  
              if self.split == 'train' and not self.scst:
                  explaination = explaination.translate(self.transtab).strip()
                  print("explaination.translate(self.transtab).strip()", explaination.translate(self.transtab).strip())
                  explaination_token_list = explaination.strip().split(" ")
                  tgt_explain = ' '.join(explaination_token_list[:self.max_tgt_length])
              else:
                  explaination = ' '.join(explaination.strip().split(" ")[:self.max_tgt_length])
                  explain_list = [explain.translate(self.transtab).strip() for explain in explaination.strip().split('&&')]
                  tgt_explain = '&&'.join(explain_list)
              print("explaination", explaination)
              print("tgt_explain", tgt_explain)
          assert False
          src_text = user_behavior + \
                     " how to persuade the user to buy the item?"
          # print("src_text", src_text.split(" ")[0:320])
          assert len(src_text.split(" ")) <= self.max_src_length
          src_item = self.encode_text(src_text)
          tgt_item = self.encode_text(" {}".format(tgt_explain))
  
          src_item = torch.cat([self.bos_item, src_item, self.eos_item])
          target_item = torch.cat([tgt_item, self.eos_item])
          prev_output_item = torch.cat([self.bos_item, tgt_item])
          example = {
              "id": uniq_id,
              "source": src_item,
              # "patch_image": patch_image,
              # "patch_mask": patch_mask,
              "target": target_item,
              "prev_output_tokens": prev_output_item
          }
          return example
  
apis: numpy.array torch.cat