content: import datetime
  import logging
  import time
  
  import numpy as np
  import torch
  import torch.distributed as dist
  import torch.nn.functional as F
  from einops import rearrange
  
  from models.criterions import get_sim
  from utils.basic_utils import MetricLogger
  from utils.distributed import get_rank, get_world_size
  
  logger = logging.getLogger(__name__)
  
  
  def extract_text_feats(texts, max_txt_l, tokenizer, model, device, return_ids=False):
      num_text = len(texts)
      text_bs = 256
      text_feats = []
      text_atts = []
      
      if return_ids:
          text_ids = []
  
      for i in range(0, num_text, text_bs):
          text = texts[i : min(num_text, i + text_bs)]
          text_input = tokenizer(
              text,
              padding="max_length",
              truncation=True,
              max_length=max_txt_l,
              return_tensors="pt",
          ).to(device) # NOTE not need to cast
  
          text_feat = model.encode_text(text_input)[0]
          text_feats.append(text_feat)
          text_atts.append(text_input.attention_mask)
          if return_ids:
              text_ids.append(text_input.input_ids)
  
      text_feats = torch.cat(text_feats, dim=0)
      text_atts = torch.cat(text_atts, dim=0)
      if return_ids:
          text_ids = torch.cat(text_ids, dim=0)
          return text_feats, text_atts, text_ids
      else:
          return text_feats, text_atts
  
  def extract_vision_feats(data_loader, model, device, config):
      if config.use_half_precision: 
          if config.get('use_bf16', False):
              cast_dtype = torch.bfloat16
          else:
              cast_dtype = torch.float16
      else:
          cast_dtype = None
  
      image_feats_all = []
      pooled_image_feats_all = []
      metric_logger = MetricLogger(delimiter="  ")
      header = "extracting image feats"
      iterator = metric_logger.log_every(data_loader, 100, header)
      for image, img_id in iterator:
          image = image.to(device, dtype=cast_dtype, non_blocking=True)
          image_feat, pooled_image_feat = model.encode_vision(image, test=True)
          if len(pooled_image_feat.shape) == 2:
              pooled_image_feat = pooled_image_feat.unsqueeze(1)  # make av_fusion happy
          if config.evaluation.eval_frame_ensemble == "concat": 
              if len(image_feat.shape) == 4:
                  image_feat = rearrange(image_feat, "b t l c -> b (t l) c").contiguous()
              image_feat = image_feat.unsqueeze(1)  # (bsz, 1, #frm*L, d)
          else:
              assert config.video_input.num_frames == 1, "only support single-frame"
              assert config.evaluation.eval_frame_ensemble in ["mean", "max", "lse"]
          if config.evaluation.eval_offload:
              image_feats_all.append(image_feat.cpu())
              pooled_image_feats_all.append(pooled_image_feat.cpu())
          else:
              image_feats_all.append(image_feat)
              pooled_image_feats_all.append(pooled_image_feat)
  
      image_feats_all = torch.cat(image_feats_all, dim=0)
      pooled_image_feats_all = torch.cat(pooled_image_feats_all, dim=0)
  
      return image_feats_all, pooled_image_feats_all
  
  def extract_audio_feats(data_loader, model, device, config):
      if config.use_half_precision: 
          if config.get('use_bf16', False):
              cast_dtype = torch.bfloat16
          else:
              cast_dtype = torch.float16
      else:
          cast_dtype = None
  
      audio_feats_all = []
      pooled_audio_feats_all = []
      metric_logger = MetricLogger(delimiter="  ")
      header = "extracting audio feats"
      iterator = metric_logger.log_every(data_loader, 100, header)
      for audio, _ in iterator:
          audio = audio.to(device, dtype=cast_dtype, non_blocking=True)
          audio_feat, pooled_audio_feat = model.encode_audio(audio, test=True)
          audio_feat = audio_feat.unsqueeze(1)  # make deep_fusion happy
          pooled_audio_feat = pooled_audio_feat.unsqueeze(1)
          if config.evaluation.eval_offload:
              audio_feats_all.append(audio_feat.cpu())
              pooled_audio_feats_all.append(pooled_audio_feat.cpu())
          else:
              audio_feats_all.append(audio_feat)
              pooled_audio_feats_all.append(pooled_audio_feat)
  
      audio_feats_all = torch.cat(audio_feats_all, dim=0)
  
      pooled_audio_feats_all = torch.cat(pooled_audio_feats_all, dim=0)
      return audio_feats_all, pooled_audio_feats_all
  
  def extract_audio_vision_feats(data_loader, model, device, config):
      if config.use_half_precision: 
          if config.get('use_bf16', False):
              cast_dtype = torch.bfloat16
          else:
              cast_dtype = torch.float16
      else:
          cast_dtype = None
  
      audio_feats_all = []
      pooled_audio_feats_all = []
      image_feats_all = []
      pooled_image_feats_all = []
      metric_logger = MetricLogger(delimiter="  ")
      header = "extracting audio and vision feats"
      iterator = metric_logger.log_every(data_loader, 100, header)
      for media, _ in iterator:
          audio = media[0]
          image = media[1]
          audio = audio.to(device, dtype=cast_dtype, non_blocking=True)
          image = image.to(device, dtype=cast_dtype, non_blocking=True)
          audio_feat, pooled_audio_feat = model.encode_audio(audio, test=True)
          audio_feat = audio_feat.unsqueeze(1)  # make deep_fusion happy
          pooled_audio_feat = pooled_audio_feat.unsqueeze(1)
          image_feat, pooled_image_feat = model.encode_vision(image, test=True)
          if len(pooled_image_feat.shape) == 2:
              pooled_image_feat = pooled_image_feat.unsqueeze(1)  # make av_fusion happy
          if config.evaluation.eval_frame_ensemble == "concat":
              if len(image_feat.shape) == 4:
                  image_feat = rearrange(image_feat, "b t l c -> b (t l) c").contiguous()
              image_feat = image_feat.unsqueeze(1)  # (bsz, 1, #frm*L, d)
          else:
              assert config.video_input.num_frames == 1, "only support single-frame"
              assert config.evaluation.eval_frame_ensemble in ["mean", "max", "lse"]
          if config.evaluation.eval_offload:
              audio_feats_all.append(audio_feat.cpu())
              pooled_audio_feats_all.append(pooled_audio_feat.cpu())
              image_feats_all.append(image_feat.cpu())
              pooled_image_feats_all.append(pooled_image_feat.cpu())
          else:
              audio_feats_all.append(audio_feat)
              pooled_audio_feats_all.append(pooled_audio_feat)
              image_feats_all.append(image_feat)
              pooled_image_feats_all.append(pooled_image_feat)
  
      audio_feats_all = torch.cat(audio_feats_all, dim=0)
      pooled_audio_feats_all = torch.cat(pooled_audio_feats_all, dim=0)
      image_feats_all = torch.cat(image_feats_all, dim=0)
      pooled_image_feats_all = torch.cat(pooled_image_feats_all, dim=0)
  
      return audio_feats_all, pooled_audio_feats_all, image_feats_all, pooled_image_feats_all
  
  
  @torch.no_grad()
  def evaluation_wrapper(model, data_loader, tokenizer, device, config, prefix=""):
      amp_eval_enabled = config.use_half_precision and not (hasattr(config, "deepspeed") and config.deepspeed.enable)
      logger.info(f"Begin to eval, model_without_ddp.dtype={model.dtype if hasattr(model, 'dtype')  else None}, amp_eval_enabled={amp_eval_enabled}, dtype={torch.bfloat16 if config.get('use_bf16', False) else torch.float16}")
      with torch.cuda.amp.autocast(enabled=amp_eval_enabled, dtype=torch.bfloat16 if config.get('use_bf16', False) else torch.float16):
          i2t_match, t2i_match = None, None
          if "qformer" in config.model.model_cls.lower():
              i2t_match, t2i_match, i2t_sim, t2i_sim, i2t_dsl, t2i_dsl = evaluation_qformer(
                  model, data_loader, tokenizer, device, config
              )
          elif "blip" in config.model.model_cls.lower():
              raise NotImplementedError
          elif "clip" in config.model.model_cls.lower() or 'coca' in config.model.model_cls.lower():
              # raise NotImplementedError
              i2t_sim, t2i_sim, i2t_dsl, t2i_dsl = evaluation_clip(
                  model, data_loader, tokenizer, device, config
              )
          else:
              i2t_match, t2i_match, i2t_sim, t2i_sim, i2t_dsl, t2i_dsl = evaluation(
                  model, data_loader, tokenizer, device, config
              )
  
          if hasattr(data_loader.dataset, "num_prompts"):
              np = data_loader.dataset.num_prompts
              logger.info(f"Using {np} prompts, we need reshape and mean!!!")
              nt = len(data_loader.dataset.text) // np
              if i2t_match is not None:
                  i2t_match = i2t_match.reshape((i2t_match.shape[0], nt, np)).mean(axis=-1)
                  t2i_match = t2i_match.reshape((nt, np, t2i_match.shape[1])).mean(axis=1)
              i2t_sim = i2t_sim.reshape((i2t_sim.shape[0], nt, np)).mean(axis=-1)
              t2i_sim = t2i_sim.reshape((nt, np, t2i_sim.shape[1])).mean(axis=1)
              i2t_dsl = i2t_dsl.reshape((i2t_dsl.shape[0], nt, np)).mean(axis=-1)
              t2i_dsl = t2i_dsl.reshape((nt, np, t2i_dsl.shape[1])).mean(axis=1)
  
      score_pairs = [
          (prefix + "_sim", i2t_sim, t2i_sim),
          (prefix + "_dsl", i2t_dsl, t2i_dsl),
      ]
      if i2t_match is not None:
          if config.evaluation.get('use_dsl_for_match', False):
              score_pairs.append((prefix + "_match (use_dsl)", i2t_match, t2i_match))
          else:
              score_pairs.append((prefix + "_match", i2t_match, t2i_match))
  
      res = dict()
      for name, i2t, t2i in score_pairs:
          if i2t is not None:
              txt2img_ids = data_loader.dataset.txt2img
              img2txt_ids = data_loader.dataset.img2txt
              res[name] = itm_eval(i2t, t2i, txt2img_ids, img2txt_ids)
      return res
  
  
  @torch.no_grad()
  def evaluation(model, data_loader, tokenizer, device, config):
      model.eval()
  
      use_dsl_for_match = config.evaluation.get('use_dsl_for_match', False)
  
      metric_logger = MetricLogger(delimiter="  ")
      header = "Evaluation:"
      dtype = torch.half if config.use_half_precision else torch.float
      media_type = data_loader.dataset.media_type
      use_subtitle = hasattr(data_loader.dataset, "use_subtitle") and data_loader.dataset.use_subtitle
      if use_subtitle:
          assert media_type in ["video", "audio_video"], f"Not support media_type: {media_type}."
          assert hasattr(data_loader.dataset, "subtitle") and data_loader.dataset.subtitle is not None, "You don't have subtitle to use."
  
      logger.info(f"Start evaluation for media_type={media_type}")
      assert media_type in ['audio', 'video', 'audio_video'], f"Not implement evaluation of {media_type}"
  
      logger.info("Computing dual encoder features...")
      start_time = time.time()
  
      # this computes all features in each GPU
      texts = data_loader.dataset.text
      # max_txt_l of eval depends on data_cofig
      max_txt_l = data_loader.dataset.max_txt_l 
  
      text_feats, text_atts = extract_text_feats(
          texts, max_txt_l, tokenizer, model, device
      )  # (bsz, Lt, d), (bsz, Lt)
  
      if use_subtitle:
          subtitle_feats, _ = extract_text_feats(
              data_loader.dataset.subtitle, max_txt_l, tokenizer, model, device
          ) # (bsz, Lt, d), (bsz, Lt)
          subtitle_proj = model.text_proj(subtitle_feats[:, 0]).unsqueeze(1)
          subtitle_feats = subtitle_feats.unsqueeze(1) 
          
      if media_type == 'video':
          image_feats, pooled_image_feats = extract_vision_feats(
              data_loader, model, device, config
          )  # (bsz, 1, #frm*Li, d) or (bsz, #frm, Li, d), (bsz, #frm, d)
          logger.info("Finished vision feature extraction")
          logger.info("Computing ITC scores [dot-product]")
          if config.evaluation.eval_offload:
              # image_feats = image_feats.to(device, non_blocking=True) image_feats will cause OOM!!!
              pooled_image_feats = pooled_image_feats.to(device, non_blocking=True)
  
          if use_subtitle:
              # print(subtitle_proj.shape, pooled_image_feats.shape)
              i2t_scores, t2i_scores = get_sim(
                  model.vs_fusion(torch.concat([subtitle_proj, model.vision_proj(pooled_image_feats)], dim=-1)), model.text_proj(text_feats[:, 0])
              )
          else:
              i2t_scores, t2i_scores = get_sim(
                  model.vision_proj(pooled_image_feats), model.text_proj(text_feats[:, 0])
              )
  
          if use_dsl_for_match:
              logger.info("use_dsl_for_match!!!")
              old_i2t_scores, old_t2i_scores = i2t_scores, t2i_scores
              i2t_scores = old_i2t_scores * old_i2t_scores.softmax(dim=0)
              t2i_scores = old_i2t_scores.T * old_i2t_scores.T.softmax(dim=0)
  
          num_medias = len(data_loader.dataset.image)
          
          # pooled_media_feats = pooled_image_feats
          if use_subtitle:
              media_feats = torch.concat([subtitle_feats, image_feats], dim=-2)
              if hasattr(model, "vstm_head"):
                  match_head = model.vstm_head
              else:
                  match_head = None
          else:
              media_feats = image_feats
              if hasattr(model, "itm_head"):
                  match_head = model.itm_head
              else:
                  match_head = None
  
      elif media_type == 'audio':
          audio_feats, pooled_audio_feats = extract_audio_feats(
              data_loader, model, device, config
          )
          logger.info("Finished audio feature extraction")
          logger.info("Computing ITC scores [dot-product]")
          if config.evaluation.eval_offload:
              pooled_audio_feats = pooled_audio_feats.to(device, non_blocking=True)
  
          i2t_scores, t2i_scores = get_sim(
              model.audio_proj(pooled_audio_feats), model.text_proj(text_feats[:, 0])
          )
  
          num_medias = len(data_loader.dataset.audio)
          media_feats = audio_feats
          # pooled_media_feats = pooled_audio_feats
          if hasattr(model, "atm_head"):
              match_head = model.atm_head
          else:
              match_head = None
  
      elif media_type == 'audio_video':
          audio_feats, pooled_audio_feats, image_feats, pooled_image_feats = extract_audio_vision_feats(
              data_loader, model, device, config
          )
          logger.info("Finished audio and vision feature extraction")
  
          logger.info("Computing ITC scores [dot-product]")
          if config.evaluation.eval_offload:
              pooled_audio_feats = pooled_audio_feats.to(device, non_blocking=True)
              pooled_image_feats = pooled_image_feats.to(device, non_blocking=True)
  
          if use_subtitle:
              i2t_scores, t2i_scores = get_sim(
                  model.avs_fusion(torch.concat([model.audio_proj(pooled_audio_feats), subtitle_proj, model.vision_proj(pooled_image_feats)], dim=-1)), model.text_proj(text_feats[:, 0])
              )
          else:
              i2t_scores, t2i_scores = get_sim(
                  model.av_fusion(torch.concat([model.audio_proj(pooled_audio_feats), model.vision_proj(pooled_image_feats)], dim=-1)), model.text_proj(text_feats[:, 0])
              )
  
          num_medias = len(data_loader.dataset.image)
          if use_subtitle:
              media_feats = torch.concat([audio_feats, subtitle_feats, image_feats], dim=-2)
              # pooled_media_feats = pooled_audio_feats
              if hasattr(model, "avstm_head"):
                  match_head = model.avstm_head
              else:
                  match_head = None
          else:
              media_feats = torch.concat([audio_feats, image_feats], dim=-2)
              # pooled_media_feats = pooled_audio_feats
              if hasattr(model, "avtm_head"):
                  match_head = model.avtm_head
              else:
                  match_head = None
      else:
          raise NotImplementedError(media_type)
      
      logger.info("Computing ITC scores [dot-product], done!")
      
      if match_head is not None:
          i2t_scores_x = torch.full((num_medias, len(texts)), -100.0).to(
              device, torch.float, non_blocking=True
          )
  
          # computes only part of the scores at each GPU, gather at the end
          logger.info("Rerank dual-encoder results with cross-encoder...")
          num_tasks = get_world_size()
          rank = get_rank()
          # only uses the part associated with the raw eval set
          # compute media2text #
          step = num_medias // num_tasks + 1
          start = rank * step
          end = min(num_medias, start + step)
  
          text_encoder = model.get_text_encoder()
          iterator = metric_logger.log_every(i2t_scores[start:end], 100, header)
          logger.info(f"i2t_scores.shape {i2t_scores[start:end].shape}")
  
          # generate score for each clip, and aggregate all clip scores for a video
          n_clip_per_video = (
              media_feats.shape[1] if not config.deep_fusion else media_feats[0].shape[1]
          )
  
          assert not config.deep_fusion and n_clip_per_video == 1, f"Not implemented for config.deep_fusion={config.deep_fusion} n_clip_per_video={n_clip_per_video}"
  
          logger.info(
              f"n_clip_per_video={n_clip_per_video}, with eval_frame_ensemble={config.evaluation.eval_frame_ensemble}"
          )
  
          for i, sims in enumerate(iterator):
              k = min(len(sims), config.evaluation.k_test)
              topk_sim, topk_idx = sims.topk(k=k, dim=0)
  
              clip_scores = []
              for clip_idx in range(n_clip_per_video):
                  if config.deep_fusion:
                      encoder_output = [
                          feat[start + i, clip_idx].to(device, non_blocking=True)
                          if config.evaluation.eval_offload
                          else feat[start + i, clip_idx]
                          for feat in media_feats
                      ]
  
                  else:
                      encoder_output = (
                          media_feats[start + i, clip_idx].to(device, non_blocking=True)
                          if config.evaluation.eval_offload
                          else media_feats[start + i, clip_idx]
                      )  # (#frm*Li, d)
  
                  # new
                  bs = 32
                  # bs = config.batch_size_test.video
                  itm_embeds = []
  
                  if config.deep_fusion:
                      if len(topk_idx) % bs != 0:
                          left = len(topk_idx) % bs
                          left_encoder_output = [feat.repeat(left, 1, 1) for feat in encoder_output]
                          left_encoder_att = [
                              torch.ones(feat.size()[:-1], dtype=torch.long).to(
                                  device, non_blocking=True
                              )
                              for feat in left_encoder_output
                          ]
                      encoder_output = [feat.repeat(bs, 1, 1) for feat in encoder_output]
                      encoder_att = [
                          torch.ones(feat.size()[:-1], dtype=torch.long).to(
                              device, non_blocking=True
                          )
                          for feat in encoder_output
                      ]
                  else:
                      if len(topk_idx) % bs != 0:
                          left = len(topk_idx) % bs
                          left_encoder_output = encoder_output.repeat(left, 1, 1)  # (k=128, #frm*Li, d)
                          left_encoder_att = torch.ones(left_encoder_output.size()[:-1], dtype=torch.long).to(
                              device, non_blocking=True
                          )
                      encoder_output = encoder_output.repeat(bs, 1, 1)  # (k=128, #frm*Li, d)
                      encoder_att = torch.ones(encoder_output.size()[:-1], dtype=torch.long).to(
                          device, non_blocking=True
                      )
  
                  for j in range(0, len(topk_idx), bs):
                      if j + bs > len(topk_idx):
                          output = text_encoder(
                              encoder_embeds=text_feats[topk_idx[j:]],
                              attention_mask=text_atts[topk_idx[j:]],
                              encoder_hidden_states=left_encoder_output,
                              encoder_attention_mask=left_encoder_att,
                              return_dict=True,
                              mode="fusion",
                          )
                      else:
                          output = text_encoder(
                              encoder_embeds=text_feats[topk_idx[j : j + bs]],
                              attention_mask=text_atts[topk_idx[j : j + bs]],
                              encoder_hidden_states=encoder_output,
                              encoder_attention_mask=encoder_att,
                              return_dict=True,
                              mode="fusion",
                          )
                      batch_itm_embeds = output.last_hidden_state[:, 0]
                      itm_embeds.append(batch_itm_embeds)
                  itm_embeds = torch.cat(itm_embeds, dim=0)
                  # end new
  
                  score = match_head(itm_embeds)[:, 1]
                  clip_scores.append(score)
  
              if len(clip_scores) == 1:
                  score = clip_scores[0]
              else:
                  raise NotImplementedError(f"len(clip_scores) == {len(clip_scores)}")
  
              i2t_scores_x[start + i, topk_idx] = score.to(i2t_scores_x.dtype)
  
          # compute text2media #
          num_text = len(data_loader.dataset.text)
          t2i_scores_x = torch.full((num_text, num_medias), -100.0).to(
              device, torch.float, non_blocking=True
          )
  
          step = num_text // num_tasks + 1
          start = rank * step
          end = min(num_text, start + step)
  
          iterator = metric_logger.log_every(t2i_scores[start:end], 100, header)
          logger.info(f"t2i_scores.shape {t2i_scores[start:end].shape}")
          # generate score for each clip, and aggregate all clip scores for a video
          n_clip_per_video = (
              media_feats.shape[1] if not config.deep_fusion else media_feats[0].shape[1]
          )
          for i, sims in enumerate(iterator):
              k = min(len(sims), config.evaluation.k_test)
              topk_sim, topk_idx = sims.topk(k=k, dim=0)
  
              clip_scores = []
              for clip_idx in range(n_clip_per_video):
                  # new
                  bs = 32
                  # bs = config.batch_size_test.video
                  itm_embeds = []
                  for j in range(0, len(topk_idx), bs):
  
                      if config.deep_fusion:
                          encoder_output = [
                              feat[topk_idx[j : j + bs].cpu(), clip_idx].to(device, non_blocking=True)
                              if config.evaluation.eval_offload
                              else feat[topk_idx[j : j + bs], clip_idx]
                              for feat in media_feats
                          ]
                          encoder_att = [
                              torch.ones(feat.size()[:-1], dtype=torch.long).to(
                                  device, non_blocking=True
                              )
                              for feat in encoder_output
                          ]
                      else:
                          encoder_output = (
                              media_feats[topk_idx[j : j + bs].cpu(), clip_idx].to(
                                  device, non_blocking=True
                              )
                              if config.evaluation.eval_offload
                              else media_feats[topk_idx[j : j + bs], clip_idx]
                          )
                          encoder_att = torch.ones(encoder_output.size()[:-1], dtype=torch.long).to(
                              device, non_blocking=True
                          )
  
                      repeat_n = (
                          encoder_output.shape[0]
                          if not config.deep_fusion
                          else encoder_output[0].shape[0]
                      )
                      output = text_encoder(
                          encoder_embeds=text_feats[start + i].repeat(repeat_n, 1, 1),
                          attention_mask=text_atts[start + i].repeat(repeat_n, 1),
                          encoder_hidden_states=encoder_output,
                          encoder_attention_mask=encoder_att,
                          return_dict=True,
                          mode="fusion",
                      )
  
                      batch_itm_embeds = output.last_hidden_state[:, 0]
                      itm_embeds.append(batch_itm_embeds)
  
                  itm_embeds = torch.cat(itm_embeds, dim=0)
                  # end new
  
                  score = match_head(itm_embeds)[:, 1]
                  clip_scores.append(score)
  
              if len(clip_scores) == 1:
                  score = clip_scores[0]
              else:
                  raise NotImplementedError(f"len(clip_scores) == {len(clip_scores)}")
  
              t2i_scores_x[start + i, topk_idx] = score.to(t2i_scores_x.dtype)
  
          logger.info("Compute over!!!")
          if config.distributed:
              logger.info("Gather across GPUs!!!")
              # gather across GPUs
              dist.barrier()
              logger.info("dist.barrier()!!!")
              dist.all_reduce(i2t_scores_x, op=dist.ReduceOp.SUM)
              logger.info("dist.all_reduce(i2t_scores_x, op=dist.ReduceOp.SUM) over!!!")
              dist.all_reduce(t2i_scores_x, op=dist.ReduceOp.SUM)
              logger.info("dist.all_reduce(t2i_scores_x, op=dist.ReduceOp.SUM) over!!!")
          
          if use_dsl_for_match:
              i2t_scores_dsl = i2t_scores
              i2t_scores_dsl_T = t2i_scores
              i2t_scores = old_i2t_scores
              t2i_scores = old_t2i_scores
          else:
              i2t_scores_dsl = i2t_scores.float() * i2t_scores.float().softmax(dim=0)
              i2t_scores_dsl_T = i2t_scores.float().T * i2t_scores.float().T.softmax(dim=0)
      else:
          i2t_scores_dsl = i2t_scores.float() * i2t_scores.float().softmax(dim=0)
          i2t_scores_dsl_T = i2t_scores.float().T * i2t_scores.float().T.softmax(dim=0)
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      logger.info(f"Evaluation time {total_time_str}")
  
  
      if match_head is not None:
          return (
              i2t_scores_x.softmax(dim=1).cpu().float().numpy(),
              t2i_scores_x.softmax(dim=1).cpu().float().numpy(),
              i2t_scores.softmax(dim=1).cpu().float().numpy(),
              i2t_scores.T.softmax(dim=1).cpu().float().numpy(),
              i2t_scores_dsl.softmax(dim=1).cpu().float().numpy(),
              i2t_scores_dsl_T.softmax(dim=1).cpu().float().numpy()
          )
      else:
          return (
              None,
              None,
              i2t_scores.softmax(dim=1).cpu().float().numpy(),
              i2t_scores.T.softmax(dim=1).cpu().float().numpy(),
              i2t_scores_dsl.softmax(dim=1).cpu().float().numpy(),
              i2t_scores_dsl_T.softmax(dim=1).cpu().float().numpy()
          )
      
  
  @torch.no_grad()
  def evaluation_simple(model, data_loader, tokenizer, device, config):
      model.eval()
  
      metric_logger = MetricLogger(delimiter="  ")
      header = "Evaluation:"
      media_type = data_loader.dataset.media_type
  
      logger.info(f"Start evaluation for media_type={media_type}")
      assert media_type in ['video'], f"Not implement evaluation of {media_type}"
  
      logger.info("Computing dual encoder features...")
      start_time = time.time()
  
      # this computes all features in each GPU
      texts = data_loader.dataset.text
      # max_txt_l of eval depends on data_cofig
      max_txt_l = data_loader.dataset.max_txt_l 
  
      text_feats, text_atts = extract_text_feats(
          texts, max_txt_l, tokenizer, model, device
      )  # (bsz, Lt, d), (bsz, Lt)
  
        
      if media_type == 'video':
          image_feats, pooled_image_feats = extract_vision_feats(
              data_loader, model, device, config
          )  # (bsz, 1, #frm*Li, d) or (bsz, #frm, Li, d), (bsz, #frm, d)
          logger.info("Finished vision feature extraction")
          logger.info("Computing ITC scores [dot-product]")
          if config.evaluation.eval_offload:
              # image_feats = image_feats.to(device, non_blocking=True) image_feats will cause OOM!!!
              pooled_image_feats = pooled_image_feats.to(device, non_blocking=True)
  
          i2t_scores, t2i_scores = get_sim(
              model.vision_proj(pooled_image_feats), model.text_proj(text_feats[:, 0])
          )
  
          num_medias = len(data_loader.dataset.image)
          
          media_feats = image_feats
          if hasattr(model, "itm_head"):
              match_head = model.itm_head
          else:
              match_head = None
  
      else:
          raise NotImplementedError(media_type)
      
      logger.info("Computing ITC scores [dot-product], done!")
      
      if match_head is not None:
          i2t_scores_x = torch.full((num_medias, len(texts)), -100.0).to(
              device, torch.float, non_blocking=True
          )
  
          # computes only part of the scores at each GPU, gather at the end
          logger.info("Rerank dual-encoder results with cross-encoder...")
          num_tasks = get_world_size()
          rank = get_rank()
          # only uses the part associated with the raw eval set
          # compute media2text #
          step = num_medias // num_tasks + 1
          start = rank * step
          end = min(num_medias, start + step)
  
          text_encoder = model.get_text_encoder()
          iterator = metric_logger.log_every(i2t_scores[start:end], 100, header)
          logger.info(f"i2t_scores.shape {i2t_scores[start:end].shape}")
  
          # generate score for each clip, and aggregate all clip scores for a video
          n_clip_per_video = (
              media_feats.shape[1] if not config.deep_fusion else media_feats[0].shape[1]
          )
  
          assert not config.deep_fusion and n_clip_per_video == 1, f"Not implemented for config.deep_fusion={config.deep_fusion} n_clip_per_video={n_clip_per_video}"
  
          logger.info(
              f"n_clip_per_video={n_clip_per_video}, with eval_frame_ensemble={config.evaluation.eval_frame_ensemble}"
          )
  
          for i, sims in enumerate(iterator):
              k = min(len(sims), config.evaluation.k_test)
              topk_sim, topk_idx = sims.topk(k=k, dim=0)
  
              clip_scores = []
              for clip_idx in range(n_clip_per_video):
                  if config.deep_fusion:
                      encoder_output = [
                          feat[start + i, clip_idx].to(device, non_blocking=True)
                          if config.evaluation.eval_offload
                          else feat[start + i, clip_idx]
                          for feat in media_feats
                      ]
  
                  else:
                      encoder_output = (
                          media_feats[start + i, clip_idx].to(device, non_blocking=True)
                          if config.evaluation.eval_offload
                          else media_feats[start + i, clip_idx]
                      )  # (#frm*Li, d)
  
                  # new
                  bs = 32
                  # bs = config.batch_size_test.video
                  itm_embeds = []
  
                  if config.deep_fusion:
                      if len(topk_idx) % bs != 0:
                          left = len(topk_idx) % bs
                          left_encoder_output = [feat.repeat(left, 1, 1) for feat in encoder_output]
                          left_encoder_att = [
                              torch.ones(feat.size()[:-1], dtype=torch.long).to(
                                  device, non_blocking=True
                              )
                              for feat in left_encoder_output
                          ]
                      encoder_output = [feat.repeat(bs, 1, 1) for feat in encoder_output]
                      encoder_att = [
                          torch.ones(feat.size()[:-1], dtype=torch.long).to(
                              device, non_blocking=True
                          )
                          for feat in encoder_output
                      ]
                  else:
                      if len(topk_idx) % bs != 0:
                          left = len(topk_idx) % bs
                          left_encoder_output = encoder_output.repeat(left, 1, 1)  # (k=128, #frm*Li, d)
                          left_encoder_att = torch.ones(left_encoder_output.size()[:-1], dtype=torch.long).to(
                              device, non_blocking=True
                          )
                      encoder_output = encoder_output.repeat(bs, 1, 1)  # (k=128, #frm*Li, d)
                      encoder_att = torch.ones(encoder_output.size()[:-1], dtype=torch.long).to(
                          device, non_blocking=True
                      )
  
                  for j in range(0, len(topk_idx), bs):
                      if j + bs > len(topk_idx):
                          output = text_encoder(
                              encoder_embeds=text_feats[topk_idx[j:]],
                              attention_mask=text_atts[topk_idx[j:]],
  