text: #!/usr/bin/env python3
  # 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.
  """
  Images and Dialogues from Image-Chat dataset.
  
  202k images, 401k utterances, over 215 different personalities.
  """
  from typing import Optional
  from parlai.core.params import ParlaiParser
  import json
  import os
  import random
  from typing import Tuple, Dict, List
  
  from parlai.core.message import Message
  from parlai.core.opt import Opt
  from parlai.core.teachers import FixedDialogTeacher
  from parlai.core.image_featurizers import ImageLoader
  from parlai.utils.typing import TShared
  from parlai.utils.io import PathManager
  from .build import build
  
  
  def _path(opt: Opt) -> Tuple[str, str, str]:
      """
      Return appropriate datapaths.
  
      :param opt:
          options
  
      :return (data path, personalities path, image_path):
          path to data, personalities, and images
      """
      build(opt)
      dt = opt['datatype'].split(':')[0]
      if dt in ['train', 'valid', 'test']:
          data_path = os.path.join(opt['datapath'], 'image_chat/{}.json'.format(dt))
  
      personalities_data_path = os.path.join(
          opt['datapath'], 'image_chat/personalities.json'
      )
      image_path = ''
      if opt.get('yfcc_path'):
          image_path = opt['yfcc_path']
      else:
          image_path = os.path.join(opt['datapath'], 'yfcc_images')
  
      return data_path, personalities_data_path, image_path
  
  
  class ImageChatTeacher(FixedDialogTeacher):
      """
      Provides the personality in the `text` field, and response in the `labels` field.
  
      To specify your own path to the YFCC100m images, please use the `--yfcc-path`
      command line argument.
      """
  
      def __init__(self, opt: Opt, shared: TShared = None):
          super().__init__(opt, shared)
          self.opt = opt
          self.image_mode = opt.get('image_mode', 'no_image_model')
          self.data_path, personalities_data_path, self.image_path = _path(opt)
          self.datatype = opt['datatype'].split(':')[0]
          self.include_personality = opt.get('include_personality')
          self.include_image = opt.get('include_image') and opt.get('load_images')
          self.num_cands = opt.get('num_cands')
          if shared and 'data' in shared:
              self.data = shared['data']
              self.personalities = shared['personalities']
              self.image_loader = shared['image_loader']
          else:
              self.image_loader = ImageLoader(opt)
              self._setup_data(self.data_path, personalities_data_path)
          self.num_exs = sum(len(d['dialog']) for d in self.data)
          self.reset()
  
      @classmethod
      def add_cmdline_args(
          cls, parser: ParlaiParser, partial_opt: Optional[Opt] = None
      ) -> ParlaiParser:
          super().add_cmdline_args(parser, partial_opt)
          agent = parser.add_argument_group('Personality-Captions arguments')
          agent.add_argument(
              '--include-personality',
              type='bool',
              default=True,
              help='Whether to provide personality to agent',
          )
          agent.add_argument(
              '--include-image',
              type='bool',
              default=True,
              help='Whether to provide image to agent',
          )
          agent.add_argument(
              '--yfcc-path',
              type=str,
              default=None,
              help='Path to yfcc images (if not downloaded '
              'via the provided download script)',
          )
          agent.add_argument(
              '--load-images',
              type='bool',
              default=True,
              help='Specify whether to load images',
          )
          agent.add_argument(
              '--num-cands',
              type=str,
              default='100',
              choices=['100', '1000'],
              help='how many candidates to provide agent',
          )
          return parser
  
      def _setup_data(self, data_path: str, personalities_data_path: str):
          """
          Load the data.
          """
          print('loading: ' + data_path)
          with PathManager.open(data_path) as f:
              self.data = json.load(f)
          with PathManager.open(personalities_data_path) as f:
              self.personalities = json.load(f)
  
      def reset(self):
          """
          Override to Reset self.example.
          """
          super().reset()
          self.example = None
  
      def num_episodes(self) -> int:
          return len(self.data)
  
      def num_examples(self) -> int:
          return self.num_exs
  
      def submit_load_request(self, image_id: str):  # type: ignore
          img_path = os.path.join(self.image_path, '{}.jpg'.format(image_id))
          self.data_loader.request_load(
              self.receive_data, self.image_loader.load, (img_path,)
          )
  
      def get(self, episode_idx: int, entry_idx: int = 0):
          data = self.data[episode_idx]
          personality, text = data['dialog'][entry_idx]
          episode_done = entry_idx == len(data['dialog']) - 1
  
          action = {
              'text': personality if self.include_personality else '',
              'image_id': data['image_hash'],
              'episode_done': episode_done,
              'labels': [text],
          }
  
          if 'candidates' in data:
              action['label_candidates'] = data['candidates'][entry_idx][self.num_cands]
  
          return action
  
      def next_example(self) -> Tuple[Message, bool]:
          """
          Returns the next example from this dataset after starting to queue up the next
          example.
  
          :return (example, epoch done):
              returns the next example as well as whether the epoch is done.
          """
          ready = None
          load_image = self.image_mode != 'no_image_model' and self.include_image
          # pull up the currently queued example
          if self.example is not None:
              # if self.image_mode != 'none' and 'image_id' in self.example:
              if load_image and 'image_id' in self.example:
                  # move the image we loaded in the background into the example
                  image = self.data_queue.get()
                  self.example['image'] = image
              ready = (self.example, self.imageEpochDone)  # type: ignore
          # get the next base example: super().next_example() calls self.get()
          self.example, self.imageEpochDone = super().next_example()
          # if self.image_mode != 'none' and 'image_id' in self.example:
          if load_image and 'image_id' in self.example:
              # load the next image in the background
              image_id = self.example['image_id']
              self.submit_load_request(image_id)
          # Try to return the previously cached example
          if ready is None:
              return self.next_example()
          else:
              return ready
  
      def share(self) -> TShared:
          shared = super().share()
          shared['data'] = self.data
          shared['image_loader'] = self.image_loader
          shared['personalities'] = self.personalities
          return shared
  
  
  class GenerationTeacher(ImageChatTeacher):
      """
      GenerationTeacher - dialogues are split into two episodes, one from
      each individual person's point of view.
  
      Used in the #dodecaDialogue task. (see https://parl.ai/projects/dodecadialogue/)
      """
  
      def __init__(self, opt: Opt, shared: TShared = None):
          if not shared:
              self.idx_to_ep = {}
          else:
              self.idx_to_ep = shared['idx_to_ep']
          self.prepend_personality = opt.get('prepend_personality', True)
          self.include_dialogue_history = opt.get('include_dialogue_history', True)
          self.category_frac = opt.get('category_frac', 0.0)
          super().__init__(opt, shared)
          self.num_eps = len(self.data) + len(
              [d for d in self.data if len(d['dialog']) > 1]
          )
  
          # Replace personalities with polarity categories ("positive/neutral" or
          # "negative"), with probability self.category_frac
          if not shared:
              category_map = get_category_map(self.personalities)
              for i, d in enumerate(self.data):
                  use_category_rand = random.random()
                  if use_category_rand < self.category_frac:
                      self.data[i]['dialog'] = [
                          [category_map[personality], label]
                          for personality, label in d['dialog']
                      ]
  
      @classmethod
      def add_cmdline_args(
          cls, parser: ParlaiParser, partial_opt: Optional[Opt] = None
      ) -> ParlaiParser:
          super().add_cmdline_args(parser, partial_opt=partial_opt)
          agent = parser.add_argument_group('generation teacher arguments')
          agent.add_argument(
              '--prepend-personality',
              type='bool',
              default=True,
              help='if true, always prepend first turn text with the personality',
          )
          agent.add_argument(
              '--include-dialogue-history',
              type='bool',
              default=True,
              help='if false, remove the dialogue history',
          )
          agent.add_argument(
              '--category-frac',
              type=float,
              default=0.0,
              help='Fraction of the time to replace the personality with its polarity category ("positive/neutral" or "negative")',
          )
          return parser
  
      def num_episodes(self) -> int:
          return self.num_eps
  
      def get(self, episode_idx: int, entry_idx: int = 0):
          entry_idx *= 2
          first_turn = entry_idx == 0
          if episode_idx >= len(self.data):
              # Second time through dataset
              data = self.data[self.idx_to_ep[episode_idx]]
              entry_idx += 1
          else:
              data = self.data[episode_idx]
  
          personality, label = data['dialog'][entry_idx]
          if not self.include_personality:
              personality = ''
  
          if entry_idx > 0:
              _, text = data['dialog'][entry_idx - 1]
              if not self.include_dialogue_history:
                  text = ''
              if first_turn and self.prepend_personality and self.include_personality:
                  text = '\n'.join([personality, text])
          elif self.prepend_personality and self.include_personality:
              text = personality
          else:
              text = ''
  
          episode_done = entry_idx >= len(data['dialog']) - 2
  
          action = {
              'text': text,
              'personality': personality,
              'image_id': data['image_hash'],
              'episode_done': episode_done,
              'labels': [label],
          }
  
          if "candidates" in data:
              action['label_candidates'] = data['candidates'][entry_idx][self.num_cands]
  
          return action
  
      def _setup_data(self, data_path: str, personalities_data_path: str):
          super()._setup_data(data_path, personalities_data_path)
          ep_idx = len(self.data)
          for i, d in enumerate(self.data):
              if len(d['dialog']) > 1:
                  self.idx_to_ep[ep_idx] = i
                  ep_idx += 1
  
      def share(self):
          shared = super().share()
          shared['idx_to_ep'] = self.idx_to_ep
          return shared
  
  
  class ImageChatTestTeacher(ImageChatTeacher):
      """
      Test ImageChat teacher for ensuring pretrained model does not break.
      """
  
      def _setup_data(self, data_path, personalities_data_path):
          super()._setup_data(data_path, personalities_data_path)
          from parlai.zoo.image_chat.transresnet_multimodal import download
  
          download(self.opt['datapath'])
          image_features_path = os.path.join(
              self.opt['datapath'],
              'models/image_chat/transresnet_multimodal/test_image_feats',
          )
          import torch
  
          with PathManager.open(image_features_path, 'rb') as f:
              self.image_features = torch.load(f)
  
      def reset(self):
          """
          Reset teacher.
          """
          super().reset()
          self.example = None
  
      def num_episodes(self):
          """
          Return number of episodes.
          """
          return len(self.image_features)
  
      def num_examples(self):
          """
          Return number of examples.
          """
          return len(self.image_features)
  
      def get(self, episode_idx, entry_idx=0):
          """
          Get an example.
  
          :param episode_idx:
              index of episode in self.data
          :param entry_idx:
              optional, which entry in the episode to get
  
          :return:
              an example
          """
          data = self.data[episode_idx]
          personality, text = data['dialog'][entry_idx]
          episode_done = entry_idx == len(data['dialog']) - 1
  
          action = {
              'text': personality if self.include_personality else '',
              'image': self.image_features[data['image_hash']],
              'episode_done': episode_done,
              'labels': [text],
          }
  
          if 'candidates' in data:
              action['label_candidates'] = data['candidates'][entry_idx][self.num_cands]
  
          return action
  
  
  class DefaultTeacher(ImageChatTeacher):
      pass
  
  
  def get_category_map(personalities: Dict[str, List[str]]) -> Dict[str, str]:
      """
      Map personalities to polarity categories: "positive/neutral" and "negative".
  
      Given a dictionary mapping Image-Chat categories (positive/neutral/negative) to
      personalities, return a dictionary mapping each personality to its category.
      Categories are merged into only two buckets: "positive/neutral", for personalities
      that are more likely to be safe to use, and "negative". Add in rare personalities.
      """
  
      category_map = {
          personality: _get_final_category(category)
          for category, personalities in personalities.items()
          for personality in personalities
      }
      category_map['Crude'] = _get_final_category('negative')
      category_map['Earnest'] = _get_final_category('positive')
      # These personalities occasionally appear but are not in personalities
      return category_map
  
  
  def _get_final_category(category: str) -> str:
      """
      Given the input raw category label, return the final one.
      """
      if category in ['positive', 'neutral']:
          return 'positive/neutral'
      elif category == 'negative':
          return 'negative'
      else:
          raise ValueError(f'Category "{category}" unrecognized!')
  