init

added_tokens.json

@@ -0,0 +1,9 @@
+{
+  "<box_begin>": 32000,
+  "<boxes>": 32003,
+  "<temp>": 32002,
+  "<time_begin>": 32001,
+  "<track_begin>": 32004,
+  "<track_box>": 32006,
+  "<tracking>": 32005
+}

config.json

@@ -0,0 +1,60 @@
+{
+  "_name_or_path": "",
+  "architectures": [
+    "MultiModalLLM_PT"
+  ],
+  "auto_map": {
+    "AutoConfig": "model_config.VideoChatEConfig",
+    "AutoModel": "modeling_videochate.MultiModalLLM_PT"
+  },
+  "model_config": {
+    "bridge": {
+      "extra_num_query_token": 64,
+      "name": "qformer",
+      "num_query_token": 32,
+      "qformer_attention_probs_dropout_prob": 0.1,
+      "qformer_drop_path_rate": 0.2,
+      "qformer_hidden_dropout_prob": 0.1
+    },
+    "freeze_bridge": false,
+    "freeze_llm": false,
+    "freeze_vision_encoder": false,
+    "llm": {
+      "lora_alpha": 32,
+      "lora_dropout": 0.1,
+      "lora_r": 16,
+      "name": "mistral_7b",
+      "pretrained_llm_path": "mistralai/Mistral-7B-Instruct-v0.3",
+      "use_lora": true,
+      "hidden_size": 4096
+    },
+    "loss": {
+      "use_vision_regression_loss": false
+    },
+    "pretrained_paths": {},
+    
+    "vision_encoder": {
+        "name":"vit_l14",
+        "img_size":224, 
+        "patch_size":16, 
+        "d_model":1024,
+        "encoder_embed_dim":1024,
+        "encoder_depth":24,
+        "encoder_num_heads":16, 
+        "drop_path_rate": 0.0, 
+        "num_frames":16,
+        "tubelet_size":1,
+        "use_checkpoint":false,
+        "checkpoint_num":0,
+        "return_index":-2,
+        "vit_add_ln":true,
+        "pretrained": null
+    }
+  },
+  "torch_dtype": "float32",
+  "transformers_version": "4.38.0",
+  "use_flash_attention": true,
+  "use_cache": true,
+  "build_decoder":true,
+  "hidden_size": 4096
+}

model-00001-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c8640081ef34803134daaa9c4a69693b680bce89b3dcc66ed22df3a26d97a330
+size 4995778624

model-00002-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8907fc1447bbaac12ba0e687cd25a19810d2717fe7e73dad57f70f13532c086c
+size 4945367960

model-00003-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5640bcb11fc22f6a7604176b29807f56a856126ed62a97a9e99585eb313e4f93
+size 4945392936

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d0360c7932a26ce651a78f5dc26e7b6735ff0bc019fb8adce231fea4375d4fe8
+size 1316534788

model_config.py

@@ -0,0 +1,24 @@
+import copy
+import re, ast
+from transformers import AutoConfig, LlamaConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from easydict import EasyDict as MyEasyDict
+from importlib import import_module
+import os.path as osp
+import argparse
+import json
+from copy import deepcopy
+import sys
+
+
+class VideoChatEConfig(PretrainedConfig):
+    model_type = 'VideoChatE'
+
+    def __init__(
+            self,
+            model_config=None,
+            **kwargs):
+        super().__init__(**kwargs)
+        self.model_config  = MyEasyDict(model_config)

modeling_base.py

@@ -0,0 +1,387 @@
+import io
+import logging
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import MSELoss
+from transformers.modeling_outputs import (
+    CausalLMOutputWithPast,
+)
+from typing import List, Optional, Tuple, Union
+from transformers import LlamaForCausalLM
+
+from torch.cuda.amp import autocast as autocast
+
+from .modeling_vit import build_vit
+from .modeling_qformer  import build_qformer
+from .model_config import VideoChatEConfig
+logger = logging.getLogger(__name__)
+
+from transformers import LlamaTokenizer,AutoTokenizer,AutoModel,AutoModelForCausalLM,AutoProcessor
+from transformers import AutoConfig, PreTrainedModel
+
+import os
+import sys
+
+
+try:
+    from third_party.sam2.build_sam import build_sam2_video_predictor
+    from third_party.cgdetr.cg_detr.model import build_cgdetr_model
+except:
+    print("can not import sam2 and cg-detr, install them first.")
+
+DEFAULT_IMG_TOKEN = "[IMG]"
+DEFAULT_IMG_END_TOKEN = "[/IMG]"
+
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_VIDEO_TOKEN = "[VIDEO]"
+
+IMG_TOKEN = "[<IMG_PLH>]"
+VID_TOKEN = "[<VID_PLH>]"
+
+BOX_START = '<box_begin>'
+# BOX_END = '<box_end>'
+ATBOXES_PLACEHOLDER = '<box_begin><boxes>'
+# ATBOXES_PLACEHOLDER = '<box_begin>'
+BOXES_PLACEHOLDER = '<boxes>'
+EXPR_PLACEHOLDER = '<expr>'
+QUESTION_PLACEHOLDER = '<question>'
+TIME_START = '<time_begin>'
+# TIME_END = '<time_end>'
+TIME_PLACEHOLDER = '<temp>'
+ATTEMP_PLACEHOLDER = TIME_START + TIME_PLACEHOLDER
+# ATTEMP_PLACEHOLDER = TIME_START
+TRACK_START='<track_begin>'
+TRACK_PLACEHOLDER = '<tracking>'
+TRACK_START_BOX = '<track_box>'
+ATTRACK_PLACEHOLDER = TRACK_START + TRACK_PLACEHOLDER
+need_template_list = ['REC', 'flickr', 'tracking', 'tracking2', 'tracking3', 'tracking4'] 
+
+load_image_list = ['image', 'REC', 'flickr']
+load_video_list = ['video', 'TVG', 'tracking', 'tracking2','tracking3', 'tracking4', 'TVG+HL']
+special_tokens = [BOX_START, TIME_START, TIME_PLACEHOLDER, BOXES_PLACEHOLDER, TRACK_START, TRACK_PLACEHOLDER, TRACK_START_BOX]
+
+def disabled_train(self, mode=True):
+    """Overwrite model.train with this function to make sure train/eval mode
+    does not change anymore."""
+    return self
+
+
+def freeze_module(module):
+    for _, param in module.named_parameters():
+        param.requires_grad = False
+    module = module.eval()
+    module.train = disabled_train
+    return module
+
+
+class LLMConfig(AutoConfig):
+    model_type = "20b"
+
+
+class BaseMLLM(PreTrainedModel):
+    config_class = VideoChatEConfig
+    def __init__(self, config,_tokenizer=None):
+        # super().__init__(config)
+        self.model_config = config.model_config
+        self.tokenizer = _tokenizer
+        
+        config.cg_opt = None
+        config.model_config = None
+        config.model_tokenizer = None
+        super().__init__(config)
+        self.build_vision_encoder()
+        self.build_llm()
+        self.build_bridge()
+        self.build_loss()
+        
+        self.load_pretrained_weights()
+        try:
+            if config.build_decoder:
+                self.cg_opt = config.cg_opt
+                self.build_bbox_decoder()
+                self.build_sam()
+                self.build_CGDETR()
+        except:
+            print("please install cgdetr and sam2 first")
+        logger.info(f'Length of tokenizer and resize embedding: {len(self.tokenizer)}')
+
+    
+    def build_vision_encoder(self):
+        if 'internvideo2' in self.model_config.vision_encoder.name.lower():
+            encoder_name = self.model_config.vision_encoder.name
+            logger.info(f"Build vision_encoder: {encoder_name}")
+            if encoder_name == 'internvideo2-1B':
+                self.vision_encoder = pretrain_internvideo2_giant_patch14_224_clean(self.model_config)
+
+            else:
+                raise ValueError(f"Not implemented: {encoder_name}")
+        elif 'vit' in self.model_config.vision_encoder.name.lower():
+            self.vision_encoder = build_vit(self.model_config)
+        else:
+            raise NotImplementedError(self.model_config.vision_encoder.name)
+
+        if self.model_config.vision_encoder.vit_add_ln:
+            self.vision_layernorm = nn.LayerNorm(self.model_config.vision_encoder.encoder_embed_dim, eps=1e-12)
+        else:
+            self.vision_layernorm = nn.Identity()
+
+        self.freeze_vision_encoder = self.model_config.get("freeze_vision_encoder", False)
+
+        if self.freeze_vision_encoder:
+            logger.info("freeze vision encoder")
+            freeze_module(self.vision_encoder)
+            freeze_module(self.vision_layernorm)
+
+    def build_CGDETR(self):
+        self.cg_model, self.cg_criterion = build_cgdetr_model()
+    
+    def build_bridge(self):
+        # ViT to LM: 1792 -> 6656 NOTE 768 is qformer dim
+        self.project_up = nn.Linear(768, self.lm.config.hidden_size) # whether bias is needed?
+        # LM to ViT: 6656 -> 1792
+        self.project_down = nn.Linear(self.lm.config.hidden_size, 768)
+        
+        if 'qformer' in self.model_config.bridge.name.lower():
+            from transformers import BertTokenizer
+            self.qformer_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased", truncation_side="left")
+            self.qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"})
+            self.qformer_tokenizer.padding_side = "left"
+            if self.model_config.bridge.name == 'qformer':
+                self.qformer, self.query_tokens = build_qformer(
+                        self.model_config.bridge.num_query_token, self.model_config.vision_encoder.encoder_embed_dim,
+                        qformer_hidden_dropout_prob=self.model_config.bridge.qformer_hidden_dropout_prob,
+                        qformer_attention_probs_dropout_prob=self.model_config.bridge.qformer_attention_probs_dropout_prob,
+                        qformer_drop_path_rate=self.model_config.bridge.qformer_drop_path_rate,
+                )
+            elif self.model_config.bridge.name == 'causal_qformer':
+                self.qformer, self.query_tokens = build_causal_qformer(
+                        self.model_config.bridge.num_query_token, self.model_config.vision_encoder.encoder_embed_dim,
+                        qformer_hidden_dropout_prob=self.model_config.bridge.qformer_hidden_dropout_prob,
+                        qformer_attention_probs_dropout_prob=self.model_config.bridge.qformer_attention_probs_dropout_prob
+                )
+            self.qformer.resize_token_embeddings(len(self.qformer_tokenizer))
+            self.qformer.cls = None
+            self.extra_num_query_token = self.model_config.bridge.extra_num_query_token
+            if self.model_config.bridge.extra_num_query_token > 0:
+                logger.info(f"Add extra {self.model_config.bridge.extra_num_query_token} tokens in QFormer")
+                self.extra_query_tokens = nn.Parameter(
+                    torch.zeros(1, self.model_config.bridge.extra_num_query_token, self.query_tokens.shape[-1])
+                )
+            
+            self.freeze_bridge = self.model_config.get("freeze_bridge", False)
+            if self.freeze_bridge:
+                logger.info("freeze bridge")
+                freeze_module(self.qformer)
+                self.query_tokens.requires_grad = False
+
+    def build_llm(self):
+        self.lm_name = self.model_config.llm.name
+        if self.model_config.llm.name == "vicuna1.5_7b":
+            self.lm = LlamaForCausalLM.from_pretrained(self.model_config.llm.pretrained_llm_path)
+            self.lm.gradient_checkpointing = self.model_config.llm.get("use_llama_gradient_checkpointing", True)
+        elif self.model_config.llm.name == 'mistral_7b':
+            from transformers import AutoModelForCausalLM
+
+            config = AutoConfig.from_pretrained(
+                self.model_config.llm.pretrained_llm_path,
+                torch_dtype=torch.bfloat16,
+                # attn_implementation="flash_attention_2",
+            )
+            self.lm = AutoModelForCausalLM.from_config(config)
+        elif self.model_config.llm.name == 'internlm_20b':
+            from transformers import AutoModelForCausalLM
+            self.lm = AutoModelForCausalLM.from_pretrained(
+                self.model_config.llm.pretrained_llm_path,
+                torch_dtype=torch.bfloat16,
+                trust_remote_code=True,
+            )
+            self.lm.gradient_checkpointing = True
+            self.lm._set_gradient_checkpointing()
+        else:
+            raise NotImplementedError(self.model_config.llm.name)
+
+        num_new_tokens = len(special_tokens)
+        self.lm.resize_token_embeddings(len(self.tokenizer))
+
+        input_embeddings = self.lm.get_input_embeddings().weight.data
+        output_embeddings = self.lm.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+        self.model_config.token_at_ids = self.tokenizer.convert_tokens_to_ids([BOX_START])[0]
+        self.freeze_llm = self.model_config.get("freeze_llm", True)
+        logger.info(f'freeze_llm: {self.freeze_llm}')
+        if self.freeze_llm:
+            logger.info("freeze llm")
+            freeze_module(self.lm)
+    
+        if self.model_config.llm.use_lora:
+            self.use_lora = True
+            from peft import get_peft_model, LoraConfig, TaskType
+            logger.info("Use lora")
+            if self.model_config.llm.name == 'internlm_20b':
+                peft_config = LoraConfig(
+                    task_type=TaskType.CAUSAL_LM, inference_mode=False, 
+                    r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout,
+                    target_modules=['wqkv', 'wo', 'w1', 'w2', 'w3', 'output']
+                )
+            else:
+                peft_config = LoraConfig(
+                    task_type=TaskType.CAUSAL_LM, inference_mode=False, 
+                    r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout,
+                    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                                    "gate_proj", "up_proj", "down_proj", "lm_head"]
+                )
+                
+            self.lm = get_peft_model(self.lm, peft_config)
+            self.lm.enable_input_require_grads()
+            self.lm.print_trainable_parameters()
+
+            if self.model_config.get("freeze_lora", False):
+                logger.info("freeze lora")
+                freeze_module(self.lm)
+                self.lm.print_trainable_parameters()
+
+        else:
+            self.use_lora = False
+
+    def add_lora(self):
+        if self.model_config.llm.use_lora:
+            self.use_lora = True
+            from peft import get_peft_model, LoraConfig, TaskType
+            logger.info("Use lora")
+            if self.model_config.llm.name == 'internlm_20b':
+                peft_config = LoraConfig(
+                    task_type=TaskType.CAUSAL_LM, inference_mode=False, 
+                    r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout,
+                    target_modules=['wqkv', 'wo', 'w1', 'w2', 'w3', 'output']
+                )
+            else:
+                peft_config = LoraConfig(
+                    task_type=TaskType.CAUSAL_LM, inference_mode=False, 
+                    r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout,
+                    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                                    "gate_proj", "up_proj", "down_proj", "lm_head"]
+                )
+                
+            self.lm = get_peft_model(self.lm, peft_config)
+            self.lm.enable_input_require_grads()
+            self.lm.print_trainable_parameters()
+
+            if self.model_config.get("freeze_lora", False):
+                logger.info("freeze lora")
+                freeze_module(self.lm)
+                self.lm.print_trainable_parameters()
+
+        else:
+            self.use_lora = False
+
+    def add_tokens(self):
+        num_new_tokens = len(special_tokens)
+        self.lm.resize_token_embeddings(len(self.tokenizer))
+
+        input_embeddings = self.lm.get_input_embeddings().weight.data
+        output_embeddings = self.lm.get_output_embeddings().weight.data
+
+        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
+            dim=0, keepdim=True)
+        print(self.lm.get_input_embeddings().weight.data.shape)
+        input_embeddings[-num_new_tokens:] = input_embeddings_avg
+        output_embeddings[-num_new_tokens:] = output_embeddings_avg
+
+        self.model_config.token_at_ids = self.tokenizer.convert_tokens_to_ids([BOX_START])[0]
+
+    def build_loss(self):
+        self.use_vision_regression_loss = self.model_config.loss.get("use_vision_regression_loss", False)
+        self.use_bbox_loss = self.model_config.loss.get("add_bbox_loss", False)
+        self.use_mask_loss = self.model_config.loss.get("use_mask_loss", False)
+        self.use_temporal_loss = self.model_config.loss.get('use_temporal_loss', False)
+        if self.use_vision_regression_loss:
+            self.image_loss_fct = MSELoss()
+        
+        
+    def load_pretrained_weights(self):
+        if self.model_config.pretrained_paths.get('pretrained_vit_qformer_path', None):
+            if 'safetensor' in self.model_config.pretrained_paths.pretrained_vit_qformer_path:
+                from safetensors import safe_open
+                from safetensors.torch import save_file
+                state_dict = {}
+                with safe_open(self.model_config.pretrained_paths.pretrained_vit_qformer_path, framework="pt", device="cpu") as f:
+                    for key in f.keys():
+                        state_dict[key] = f.get_tensor(key)
+            else:
+                state_dict = torch.load(self.model_config.pretrained_paths.pretrained_vit_qformer_path, map_location="cpu")
+                if "model" in state_dict.keys():
+                    state_dict = state_dict["model"]
+                elif "module" in state_dict.keys(): 
+                    state_dict = state_dict["module"] # for deepspeed
+            self.check_temp_emb(state_dict)
+            msg = self.load_state_dict(state_dict, strict=False)
+            print('Loading vit: ', msg)
+            logger.info(f"Load ViT and QFormer from {self.model_config.pretrained_paths.pretrained_vit_qformer_path}: {msg}")
+
+        if self.model_config.pretrained_paths.get('pretrained_videochat2', None):
+            state_dict = torch.load(self.model_config.pretrained_paths.pretrained_videochat2, map_location="cpu")
+            
+            new_state_dict = {}
+            for k in state_dict.keys():
+                if 'bert.embeddings' not in k:
+                    new_state_dict[k] = state_dict[k]
+            state_dict = new_state_dict
+            # self.check_temp_emb(state_dict)
+            msg = self.load_state_dict(state_dict, strict=False)
+            print('Loading videochat2: ', msg)
+        
+
+    def check_temp_emb(self, state_dict):
+        old_num_frames = self.model_config.vision_encoder.get('origin_num_frames', None)
+        new_num_frames = self.model_config.vision_encoder.num_frames
+        if old_num_frames is not None and old_num_frames != new_num_frames:
+            logger.info(f"interpolate_pos_embed_internvideo2 to {new_num_frames} (origin_num_frames={old_num_frames})!!!")
+            a = len(state_dict)
+            interpolate_pos_embed_internvideo2_new(state_dict, self.vision_encoder, orig_t_size=4)
+            assert a == len(state_dict), state_dict.keys()
+
+    def build_bbox_decoder(self):
+        self.loc_encoder = nn.Sequential(
+            nn.Linear(4, self.model_config.llm.hidden_size // 2, dtype=torch.bfloat16),
+            nn.ReLU(),
+            nn.Linear(self.model_config.llm.hidden_size // 2, self.model_config.llm.hidden_size, dtype=torch.bfloat16),
+        )
+
+        self.loc_decoder = nn.Sequential(
+            nn.Linear(self.model_config.llm.hidden_size, self.model_config.llm.hidden_size // 2, dtype=torch.bfloat16),
+            nn.ReLU(),
+            nn.Linear(self.model_config.llm.hidden_size // 2, 4, dtype=torch.bfloat16)
+        )
+        self._initialize_bbox_weights()
+
+    def _initialize_bbox_weights(self):
+        return
+
+    def build_sam(self):
+        sam2_checkpoint = "/cpfs01/user/heyinan/checkpoints/sam2_hiera_large.pt"
+        model_cfg = "sam2_hiera_l.yaml"
+        predictor = build_sam2_video_predictor(model_cfg, sam2_checkpoint, device=self.lm.device)
+        
+        self.sam = predictor
+        freeze_module(self.sam)
+        
+
+    @property
+    def dtype(self):
+        return self.lm.dtype
+
+
+    @property
+    def device(self):
+        return self.lm.device

modeling_qformer.py

@@ -0,0 +1,1264 @@
+"""
+ * Copyright (c) 2023, salesforce.com, inc.
+ * All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ * For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+ * By Junnan Li
+ * Based on huggingface code base
+ * https://github.com/huggingface/transformers/blob/v4.15.0/src/transformers/models/bert
+"""
+import logging
+import math
+import os
+import warnings
+from dataclasses import dataclass
+from typing import Optional, Tuple, Dict, Any
+
+import torch
+from torch import Tensor, device, dtype, nn
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import torch.nn.functional as F
+
+from timm.models.layers import drop_path
+from transformers.activations import ACT2FN
+from transformers.file_utils import (
+    ModelOutput,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    MaskedLMOutput,
+    MultipleChoiceModelOutput,
+    NextSentencePredictorOutput,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from transformers.models.bert.configuration_bert import BertConfig
+
+import logging
+logger = logging.getLogger(__name__)
+
+
+class BertEmbeddings(nn.Module):
+    """Construct the embeddings from word and position embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))
+        )
+        self.position_embedding_type = getattr(
+            config, "position_embedding_type", "absolute"
+        )
+
+        self.config = config
+
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        query_embeds=None,
+        past_key_values_length=0,
+    ):
+        if input_ids is not None:
+            seq_length = input_ids.size()[1]
+        else:
+            seq_length = 0
+
+        if position_ids is None:
+            position_ids = self.position_ids[
+                :, past_key_values_length : seq_length + past_key_values_length
+            ].clone()
+
+        if input_ids is not None:
+            embeddings = self.word_embeddings(input_ids)
+            if self.position_embedding_type == "absolute":
+                position_embeddings = self.position_embeddings(position_ids)
+                embeddings = embeddings + position_embeddings
+
+            if query_embeds is not None:
+                embeddings = torch.cat((query_embeds, embeddings), dim=1)
+        else:
+            embeddings = query_embeds
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class BertSelfAttention(nn.Module):
+    def __init__(self, config, is_cross_attention):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(
+            config, "embedding_size"
+        ):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_width, self.all_head_size)
+            self.value = nn.Linear(config.encoder_width, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.position_embedding_type = getattr(
+            config, "position_embedding_type", "absolute"
+        )
+        if (
+            self.position_embedding_type == "relative_key"
+            or self.position_embedding_type == "relative_key_query"
+        ):
+            self.max_position_embeddings = config.max_position_embeddings
+            self.distance_embedding = nn.Embedding(
+                2 * config.max_position_embeddings - 1, self.attention_head_size
+            )
+        self.save_attention = False
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+
+    def get_attention_map(self):
+        return self.attention_map
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (
+            self.num_attention_heads,
+            self.attention_head_size,
+        )
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+
+        if is_cross_attention:
+            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+            attention_mask = encoder_attention_mask
+        elif past_key_value is not None:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+        else:
+            key_layer = self.transpose_for_scores(self.key(hidden_states))
+            value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+        mixed_query_layer = self.query(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+
+        past_key_value = (key_layer, value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        if (
+            self.position_embedding_type == "relative_key"
+            or self.position_embedding_type == "relative_key_query"
+        ):
+            seq_length = hidden_states.size()[1]
+            position_ids_l = torch.arange(
+                seq_length, dtype=torch.long, device=hidden_states.device
+            ).view(-1, 1)
+            position_ids_r = torch.arange(
+                seq_length, dtype=torch.long, device=hidden_states.device
+            ).view(1, -1)
+            distance = position_ids_l - position_ids_r
+            positional_embedding = self.distance_embedding(
+                distance + self.max_position_embeddings - 1
+            )
+            positional_embedding = positional_embedding.to(
+                dtype=query_layer.dtype
+            )  # fp16 compatibility
+
+            if self.position_embedding_type == "relative_key":
+                relative_position_scores = torch.einsum(
+                    "bhld,lrd->bhlr", query_layer, positional_embedding
+                )
+                attention_scores = attention_scores + relative_position_scores
+            elif self.position_embedding_type == "relative_key_query":
+                relative_position_scores_query = torch.einsum(
+                    "bhld,lrd->bhlr", query_layer, positional_embedding
+                )
+                relative_position_scores_key = torch.einsum(
+                    "bhrd,lrd->bhlr", key_layer, positional_embedding
+                )
+                attention_scores = (
+                    attention_scores
+                    + relative_position_scores_query
+                    + relative_position_scores_key
+                )
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        if is_cross_attention and self.save_attention:
+            self.save_attention_map(attention_probs)
+            attention_probs.register_hook(self.save_attn_gradients)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs_dropped = attention_probs_dropped * head_mask
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (
+            (context_layer, attention_probs) if output_attentions else (context_layer,)
+        )
+
+        outputs = outputs + (past_key_value,)
+        return outputs
+    
+    
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class BertSelfOutput(nn.Module):
+    def __init__(self, config, drop_path=0.):
+        super().__init__()
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.drop_path(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False, drop_path=0.,):
+        super().__init__()
+        self.self = BertSelfAttention(config, is_cross_attention)
+        self.output = BertSelfOutput(config, drop_path=drop_path)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads,
+            self.self.num_attention_heads,
+            self.self.attention_head_size,
+            self.pruned_heads,
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = (
+            self.self.attention_head_size * self.self.num_attention_heads
+        )
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            past_key_value,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+
+        outputs = (attention_output,) + self_outputs[
+            1:
+        ]  # add attentions if we output them
+        return outputs
+
+
+class BertIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+class BertOutput(nn.Module):
+    def __init__(self, config, drop_path=0.):
+        super().__init__()
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.drop_path(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BertLayer(nn.Module):
+    def __init__(self, config, layer_num):
+        super().__init__()
+        self.config = config
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        drop_path = config.drop_path_list[layer_num]
+        self.attention = BertAttention(config, drop_path=drop_path)
+        self.layer_num = layer_num
+        if (
+            self.config.add_cross_attention
+            and layer_num % self.config.cross_attention_freq == 0
+        ):
+            self.crossattention = BertAttention(
+                config, is_cross_attention=self.config.add_cross_attention,
+                drop_path=drop_path
+            )
+            self.has_cross_attention = True
+        else:
+            self.has_cross_attention = False
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config, drop_path=drop_path)
+
+        self.intermediate_query = BertIntermediate(config)
+        self.output_query = BertOutput(config, drop_path=drop_path)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        output_attentions=False,
+        query_length=0,
+    ):
+        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+        self_attn_past_key_value = (
+            past_key_value[:2] if past_key_value is not None else None
+        )
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+            past_key_value=self_attn_past_key_value,
+        )
+        attention_output = self_attention_outputs[0]
+        outputs = self_attention_outputs[1:-1]
+
+        present_key_value = self_attention_outputs[-1]
+
+        if query_length > 0:
+            query_attention_output = attention_output[:, :query_length, :]
+
+            if self.has_cross_attention:
+                assert (
+                    encoder_hidden_states is not None
+                ), "encoder_hidden_states must be given for cross-attention layers"
+                cross_attention_outputs = self.crossattention(
+                    query_attention_output,
+                    attention_mask,
+                    head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions=output_attentions,
+                )
+                query_attention_output = cross_attention_outputs[0]
+                outputs = (
+                    outputs + cross_attention_outputs[1:-1]
+                )  # add cross attentions if we output attention weights
+
+            layer_output = apply_chunking_to_forward(
+                self.feed_forward_chunk_query,
+                self.chunk_size_feed_forward,
+                self.seq_len_dim,
+                query_attention_output,
+            )
+            if attention_output.shape[1] > query_length:
+                layer_output_text = apply_chunking_to_forward(
+                    self.feed_forward_chunk,
+                    self.chunk_size_feed_forward,
+                    self.seq_len_dim,
+                    attention_output[:, query_length:, :],
+                )
+                layer_output = torch.cat([layer_output, layer_output_text], dim=1)
+        else:
+            layer_output = apply_chunking_to_forward(
+                self.feed_forward_chunk,
+                self.chunk_size_feed_forward,
+                self.seq_len_dim,
+                attention_output,
+            )
+        outputs = (layer_output,) + outputs
+
+        outputs = outputs + (present_key_value,)
+
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+    def feed_forward_chunk_query(self, attention_output):
+        intermediate_output = self.intermediate_query(attention_output)
+        layer_output = self.output_query(intermediate_output, attention_output)
+        return layer_output
+
+
+class BertEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList(
+            [BertLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+        query_length=0,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attentions = () if output_attentions else None
+        all_cross_attentions = (
+            () if output_attentions and self.config.add_cross_attention else None
+        )
+
+        next_decoder_cache = () if use_cache else None
+
+        for i in range(self.config.num_hidden_layers):
+            layer_module = self.layer[i]
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+            past_key_value = past_key_values[i] if past_key_values is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False) and self.training:
+
+                if use_cache:
+                    logger.warn(
+                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                    )
+                    use_cache = False
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(
+                            *inputs, past_key_value, output_attentions, query_length
+                        )
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    past_key_value,
+                    output_attentions,
+                    query_length,
+                )
+
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                next_decoder_cache += (layer_outputs[-1],)
+            if output_attentions:
+                all_self_attentions = all_self_attentions + (layer_outputs[1],)
+                all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_decoder_cache,
+                    all_hidden_states,
+                    all_self_attentions,
+                    all_cross_attentions,
+                ]
+                if v is not None
+            )
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=next_decoder_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attentions,
+            cross_attentions=all_cross_attentions,
+        )
+
+
+class BertPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = BertLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class BertPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = BertConfig
+    base_model_prefix = "bert"
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+class BertModel(BertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an :obj:`encoder_hidden_states` is then expected as an
+    input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=False):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BertEmbeddings(config)
+
+        self.encoder = BertEncoder(config)
+
+        self.pooler = BertPooler(config) if add_pooling_layer else None
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def get_extended_attention_mask(
+        self,
+        attention_mask: Tensor,
+        input_shape: Tuple[int],
+        device: device,
+        is_decoder: bool,
+        has_query: bool = False,
+    ) -> Tensor:
+        """
+        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+        Arguments:
+            attention_mask (:obj:`torch.Tensor`):
+                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+            input_shape (:obj:`Tuple[int]`):
+                The shape of the input to the model.
+            device: (:obj:`torch.device`):
+                The device of the input to the model.
+
+        Returns:
+            :obj:`torch.Tensor` The extended attention mask, with a the same dtype as :obj:`attention_mask.dtype`.
+        """
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if attention_mask.dim() == 3:
+            extended_attention_mask = attention_mask[:, None, :, :]
+        elif attention_mask.dim() == 2:
+            # Provided a padding mask of dimensions [batch_size, seq_length]
+            # - if the model is a decoder, apply a causal mask in addition to the padding mask
+            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+            if is_decoder:
+                batch_size, seq_length = input_shape
+
+                seq_ids = torch.arange(seq_length, device=device)
+                causal_mask = (
+                    seq_ids[None, None, :].repeat(batch_size, seq_length, 1)
+                    <= seq_ids[None, :, None]
+                )
+
+                # add a prefix ones mask to the causal mask
+                # causal and attention masks must have same type with pytorch version < 1.3
+                causal_mask = causal_mask.to(attention_mask.dtype)
+
+                if causal_mask.shape[1] < attention_mask.shape[1]:
+                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
+                    if has_query:  # UniLM style attention mask
+                        causal_mask = torch.cat(
+                            [
+                                torch.zeros(
+                                    (batch_size, prefix_seq_len, seq_length),
+                                    device=device,
+                                    dtype=causal_mask.dtype,
+                                ),
+                                causal_mask,
+                            ],
+                            axis=1,
+                        )
+                    causal_mask = torch.cat(
+                        [
+                            torch.ones(
+                                (batch_size, causal_mask.shape[1], prefix_seq_len),
+                                device=device,
+                                dtype=causal_mask.dtype,
+                            ),
+                            causal_mask,
+                        ],
+                        axis=-1,
+                    )
+                extended_attention_mask = (
+                    causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+                )
+            else:
+                extended_attention_mask = attention_mask[:, None, None, :]
+        else:
+            raise ValueError(
+                "Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+                    input_shape, attention_mask.shape
+                )
+            )
+
+        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+        # masked positions, this operation will create a tensor which is 0.0 for
+        # positions we want to attend and -10000.0 for masked positions.
+        # Since we are adding it to the raw scores before the softmax, this is
+        # effectively the same as removing these entirely.
+        extended_attention_mask = extended_attention_mask.to(
+            dtype=self.dtype
+        )  # fp16 compatibility
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        return extended_attention_mask
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        query_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_values=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        is_decoder=False,
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        """
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if input_ids is None:
+            assert (
+                query_embeds is not None
+            ), "You have to specify query_embeds when input_ids is None"
+
+        # past_key_values_length
+        past_key_values_length = (
+            past_key_values[0][0].shape[2] - self.config.query_length
+            if past_key_values is not None
+            else 0
+        )
+
+        query_length = query_embeds.shape[1] if query_embeds is not None else 0
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            query_embeds=query_embeds,
+            past_key_values_length=past_key_values_length,
+        )
+
+        input_shape = embedding_output.size()[:-1]
+        batch_size, seq_length = input_shape
+        device = embedding_output.device
+
+        if attention_mask is None:
+            attention_mask = torch.ones(
+                ((batch_size, seq_length + past_key_values_length)), device=device
+            )
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        if is_decoder:
+            extended_attention_mask = self.get_extended_attention_mask(
+                attention_mask,
+                input_ids.shape,
+                device,
+                is_decoder,
+                has_query=(query_embeds is not None),
+            )
+        else:
+            extended_attention_mask = self.get_extended_attention_mask(
+                attention_mask, input_shape, device, is_decoder
+            )
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if encoder_hidden_states is not None:
+            if type(encoder_hidden_states) == list:
+                encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[
+                    0
+                ].size()
+            else:
+                (
+                    encoder_batch_size,
+                    encoder_sequence_length,
+                    _,
+                ) = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+            if type(encoder_attention_mask) == list:
+                encoder_extended_attention_mask = [
+                    self.invert_attention_mask(mask) for mask in encoder_attention_mask
+                ]
+            elif encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+                encoder_extended_attention_mask = self.invert_attention_mask(
+                    encoder_attention_mask
+                )
+            else:
+                encoder_extended_attention_mask = self.invert_attention_mask(
+                    encoder_attention_mask
+                )
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            query_length=query_length,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = (
+            self.pooler(sequence_output) if self.pooler is not None else None
+        )
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+class BertLMHeadModel(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        query_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        past_key_values=None,
+        use_cache=True,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        return_logits=False,
+        is_decoder=True,
+        reduction="mean",
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in ``[0, 1]``:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are
+            ignored (masked), the loss is only computed for the tokens with labels n ``[0, ..., config.vocab_size]``
+        past_key_values (:obj:`tuple(tuple(torch.FloatTensor))` of length :obj:`config.n_layers` with each tuple having 4 tensors of shape :obj:`(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If :obj:`past_key_values` are used, the user can optionally input only the last :obj:`decoder_input_ids`
+            (those that don't have their past key value states given to this model) of shape :obj:`(batch_size, 1)`
+            instead of all :obj:`decoder_input_ids` of shape :obj:`(batch_size, sequence_length)`.
+        use_cache (:obj:`bool`, `optional`):
+            If set to :obj:`True`, :obj:`past_key_values` key value states are returned and can be used to speed up
+            decoding (see :obj:`past_key_values`).
+        Returns:
+        Example::
+            >>> from transformers import BertTokenizer, BertLMHeadModel, BertConfig
+            >>> import torch
+            >>> tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
+            >>> config = BertConfig.from_pretrained("bert-base-cased")
+            >>> model = BertLMHeadModel.from_pretrained('bert-base-cased', config=config)
+            >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+            >>> outputs = model(**inputs)
+            >>> prediction_logits = outputs.logits
+        """
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        if labels is not None:
+            use_cache = False
+        if past_key_values is not None:
+            query_embeds = None
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            query_embeds=query_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+        )
+
+        sequence_output = outputs[0]
+        if query_embeds is not None:
+            sequence_output = outputs[0][:, query_embeds.shape[1] :, :]
+
+        prediction_scores = self.cls(sequence_output)
+
+        if return_logits:
+            return prediction_scores[:, :-1, :].contiguous()
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous()
+            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=0.1)
+            lm_loss = loss_fct(
+                shifted_prediction_scores.view(-1, self.config.vocab_size),
+                labels.view(-1),
+            )
+            if reduction == "none":
+                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((lm_loss,) + output) if lm_loss is not None else output
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=lm_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self, input_ids, query_embeds, past=None, attention_mask=None, **model_kwargs
+    ):
+        # if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+        if attention_mask is None:
+            attention_mask = input_ids.new_ones(input_ids.shape)
+        query_mask = input_ids.new_ones(query_embeds.shape[:-1])
+        attention_mask = torch.cat([query_mask, attention_mask], dim=-1)
+
+        # cut decoder_input_ids if past is used
+        if past is not None:
+            input_ids = input_ids[:, -1:]
+
+        return {
+            "input_ids": input_ids,
+            "query_embeds": query_embeds,
+            "attention_mask": attention_mask,
+            "past_key_values": past,
+            "encoder_hidden_states": model_kwargs.get("encoder_hidden_states", None),
+            "encoder_attention_mask": model_kwargs.get("encoder_attention_mask", None),
+            "is_decoder": True,
+        }
+
+    def _reorder_cache(self, past, beam_idx):
+        reordered_past = ()
+        for layer_past in past:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx) for past_state in layer_past
+                ),
+            )
+        return reordered_past
+
+
+class BertForMaskedLM(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids", r"predictions.decoder.bias"]
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.cls = BertOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        position_ids=None,
+        head_mask=None,
+        query_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        return_logits=False,
+        is_decoder=False,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the masked language modeling loss. Indices should be in ``[-100, 0, ...,
+            config.vocab_size]`` (see ``input_ids`` docstring) Tokens with indices set to ``-100`` are ignored
+            (masked), the loss is only computed for the tokens with labels in ``[0, ..., config.vocab_size]``
+        """
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            query_embeds=query_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            is_decoder=is_decoder,
+        )
+
+        if query_embeds is not None:
+            sequence_output = outputs[0][:, query_embeds.shape[1] :, :]
+        prediction_scores = self.cls(sequence_output)
+
+        if return_logits:
+            return prediction_scores
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size), labels.view(-1)
+            )
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return (
+                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            )
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+def build_qformer(num_query_token, vision_width, 
+                  qformer_hidden_dropout_prob=0.1,
+                  qformer_attention_probs_dropout_prob=0.1,
+                  qformer_drop_path_rate=0.,
+                  bert_type="bert-base-uncased"
+                  ):
+    
+    encoder_config = BertConfig.from_pretrained(bert_type)
+    encoder_config.encoder_width = vision_width
+    # insert cross-attention layer every other block
+    encoder_config.add_cross_attention = True
+    encoder_config.cross_attention_freq = 2
+    encoder_config.query_length = num_query_token
+    encoder_config.hidden_dropout_prob = qformer_hidden_dropout_prob
+    encoder_config.attention_probs_dropout_prob = qformer_attention_probs_dropout_prob
+    encoder_config.drop_path_list = [x.item() for x in torch.linspace(0, qformer_drop_path_rate, encoder_config.num_hidden_layers)]
+    logger.info(f"Drop_path:{encoder_config.drop_path_list}")
+    logger.info(encoder_config)
+    Qformer = BertLMHeadModel(encoder_config)                     
+    query_tokens = nn.Parameter(
+        torch.zeros(1, num_query_token, encoder_config.hidden_size)
+    )
+    query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range)
+    return Qformer, query_tokens
+

modeling_special_token.py

@@ -0,0 +1,27 @@
+import transformers
+DEFAULT_IMG_TOKEN = "[IMG]"
+DEFAULT_IMG_END_TOKEN = "[/IMG]"
+
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_VIDEO_TOKEN = "[VIDEO]"
+
+IMG_TOKEN = "[<IMG_PLH>]"
+VID_TOKEN = "[<VID_PLH>]"
+
+BOX_START = '<box_begin>'
+ATBOXES_PLACEHOLDER = '<box_begin><boxes>'
+BOXES_PLACEHOLDER = '<boxes>'
+EXPR_PLACEHOLDER = '<expr>'
+QUESTION_PLACEHOLDER = '<question>'
+TIME_START = '<time_begin>'
+TIME_PLACEHOLDER = '<temp>'
+ATTEMP_PLACEHOLDER = TIME_START + TIME_PLACEHOLDER
+TRACK_START='<track_begin>'
+TRACK_PLACEHOLDER = '<tracking>'
+TRACK_START_BOX = '<track_box>'
+ATTRACK_PLACEHOLDER = TRACK_START + TRACK_PLACEHOLDER
+need_template_list = ['REC', 'flickr', 'tracking', 'tracking2', 'tracking3', 'tracking4'] 
+
+load_image_list = ['image', 'REC', 'flickr']
+load_video_list = ['video', 'TVG', 'tracking', 'tracking2','tracking3', 'tracking4', 'TVG+HL']
+special_tokens = [BOX_START, TIME_START, TIME_PLACEHOLDER, BOXES_PLACEHOLDER, TRACK_START, TRACK_PLACEHOLDER, TRACK_START_BOX]

modeling_videochate.py

@@ -0,0 +1,681 @@
+import io
+import logging
+import json
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import MSELoss
+from transformers.modeling_outputs import (
+    CausalLMOutputWithPast,
+)
+from typing import List, Optional, Tuple, Union
+from transformers import LlamaForCausalLM
+from transformers.modeling_outputs import (
+    CausalLMOutputWithPast,
+)
+
+from torch.cuda.amp import autocast as autocast
+import torch.nn.functional as F
+
+import numpy as np
+from .modeling_vit import  build_vit, MLP, PostProcess
+
+from .modeling_qformer import build_qformer
+from .modeling_base import BaseMLLM
+
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+logger = logging.getLogger(__name__)
+
+import pycocotools.mask as mask_util
+
+from .modeling_base import VID_TOKEN, IMG_TOKEN
+
+class MultiModalLLM_PT(BaseMLLM):
+    def __init__(
+        self,
+        config,
+        _tokenizer=None
+    ):
+        super().__init__(config=config, _tokenizer=_tokenizer)
+        self.use_clip = False
+        self.num_frames = 16
+        self.num_clips = 1
+        self.token_merge_len = 4
+
+        self.per_clip_frames = self.num_frames // self.num_clips
+        
+        print(self.config)
+        self.merge_proj = nn.Linear(
+            self.qformer.config.hidden_size*self.token_merge_len, self.config.hidden_size
+        )
+
+        if config.build_decoder:
+            self.track_embed = MLP(self.config.hidden_size, self.config.hidden_size, 3 * 256, 2, dropout=0)
+            self.track_embed_decode2 = MLP(4096, 4096, 4, 2, dropout=0)
+            self.temporal_embed = MLP(self.config.hidden_size, self.config.hidden_size, 2, 2, dropout=0.3)
+            self.action_embed = MLP(self.config.hidden_size, self.config.hidden_size, 1, 2, dropout=0.3)
+            self.postprocess = PostProcess()
+            self.track_token = nn.Parameter(torch.randn((1, 1, 4096)))
+            self.temporal_token = nn.Parameter(torch.randn((1, 1, 4096)))
+            self.box_token = nn.Parameter(torch.randn((1, 1, 4096)))
+
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        image: Optional[torch.Tensor] = None,
+        video: Optional[torch.Tensor] = None,
+        instruction = None,
+        video_idx = None,
+        image_idx = None,
+        output_boxes = None, # REC
+        input_boxes = None, # tracking inputs
+        text_input = None,
+        video_info = None, 
+        temporal_labels = None,
+        gt_masks = None,
+        sam_images = None,
+        size_hw = None,
+        path = None,
+        mask_path = None,
+        tvg_inputs = None,
+        tvg_targets = None,
+    ):  
+        if text_input is not None:
+            time_instructions = self.get_clip_time_instruct(text_input)
+        else:
+            time_instructions = None
+        text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, return_visual=False,
+                                        video_idx=video_idx, image_idx=image_idx,  instruction = instruction, 
+                                        output_boxes = output_boxes, input_boxes=input_boxes, time_instructions = time_instructions)
+        outputs = self.lm(
+            inputs_embeds=text_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        loss = outputs.loss
+        logger.info(f'llm loss:{loss}')
+
+        if output_boxes is not None and self.use_bbox_loss:
+            last_hidden_states = outputs.hidden_states[-1]
+            pred_locs = []
+            for idx in range(last_hidden_states.shape[0]):
+                loc_positions = ( (input_ids[idx].flatten() == self.tokenizer.box_token) ).nonzero().flatten()
+                selected_hidden_states = last_hidden_states[idx][loc_positions]
+                pred_locs.append(self.loc_decoder(selected_hidden_states))
+            box_loss = self.box_loss(pred_locs, output_boxes)
+            logger.info(f'box loss:{box_loss}')
+            loss += box_loss
+        
+        if (gt_masks is not None or input_boxes is not None) and self.use_mask_loss:
+            last_hidden_states = outputs.hidden_states[-1]
+            pred_masks = []
+            sam_losses = []
+            box_losses = []
+            for idx in range(last_hidden_states.shape[0]):
+                loc_positions = ( (input_ids[idx].flatten() == self.tokenizer.track_token) ).nonzero().flatten()
+                selected_hidden_states = last_hidden_states[idx][loc_positions]
+                embed_sam_boxes = self.track_embed(selected_hidden_states).reshape(1, 3, 256)
+                inference_state = self.sam.init_state_images(sam_images, size_hw[idx][0], size_hw[idx][1])
+                
+                if input_boxes is not None:
+                    gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes[idx], device = text_embeds.device) 
+                else:
+                    input_boxes = self.find_boundaries_torch(gt_masks.squeeze(0)[:,:,:1].squeeze(2).cpu()).to(text_embeds.device)
+                    gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes, device = text_embeds.device) 
+                pred_locs = [self.track_embed_decode2((selected_hidden_states))[0]]
+                target_boxes = [input_boxes[idx]]
+
+                src_boxes = pred_locs
+                loss_bbox = self.box_loss2(src_boxes, target_boxes)
+
+                loss_bbox = self.masked_loss(loss_bbox, 0)
+                box_losses.append(loss_bbox)
+                sam_losses.append( F.l1_loss(embed_sam_boxes, gt_embeds))
+            
+            logger.info(f'refering sam loss:{sam_losses}')
+            sam_losses = torch.stack(sam_losses)
+            box_losses = torch.stack(box_losses)
+            loss += torch.mean(sam_losses)
+            loss += torch.mean(box_losses)
+        
+        if tvg_inputs is not None and self.use_temporal_loss:
+            last_hidden_states = outputs.hidden_states[-1]                                               # [bsz,1024, 4096]
+            last_hidden_states = last_hidden_states.view(-1, last_hidden_states.size(-1))                # [bsz*1024, 4096]
+            loc_positions = (input_ids.flatten()==self.tokenizer.temp_token).nonzero().flatten()         # [bsz]
+            prompt_token = last_hidden_states[loc_positions]
+            prompt_token = prompt_token.view(input_ids.shape[0], -1 ,prompt_token.shape[-1])   # [bsz, 1, 4096]
+
+        
+            cg_outputs = self.cg_model(**tvg_inputs, targets=tvg_targets, prompt_token=prompt_token)
+            loss_dict = self.cg_criterion(cg_outputs, tvg_targets)
+            weight_dict = self.cg_criterion.weight_dict
+            tvg_loss = 0.05*sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
+            logger.info(f'tvg_loss:{tvg_loss}')
+            loss += tvg_loss
+
+
+        logger.info(f'all loss:{loss}')
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=outputs.logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def pad_text_embeds(
+        self,
+        input_ids: torch.LongTensor = None,
+        image: Optional[torch.Tensor] = None,
+        video: Optional[torch.Tensor] = None,
+        image_idx = None,
+        video_idx = None,
+        return_visual: bool = False,
+        instruction = None,
+        output_boxes = None, # boxes for REC
+        input_boxes = None, # boxes for tracking
+        time_instructions = None,
+    ):
+        text_embeds = self.lm.get_input_embeddings()(input_ids.long()).detach()
+        if input_boxes is not None:
+            input_boxes = input_boxes[0].to(dtype=text_embeds.dtype)
+            
+            boxes_emb = self.loc_encoder(input_boxes)
+            boxes_emb = boxes_emb.view(-1, 4096)
+            
+            text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)] = text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)] * 0 + boxes_emb.to(text_embeds.device)
+            logger.info(f'embedings:{text_embeds[input_ids == torch.full_like(input_ids, self.tokenizer.track_box_token)].shape}')
+        visual = None
+        visual_idx = None
+     
+        if image is not None:
+
+            B, T, C, H, W = image.shape
+            image = image.permute(0, 2, 1, 3, 4)
+
+            instruction = None
+            
+            prompt_image_embeds = self.encode_vision(image, instruction)
+
+            visual = prompt_image_embeds
+
+            prompt_image_embeds = self.project_up(prompt_image_embeds) # 768 -> 4096
+            prompt_image_embeds = prompt_image_embeds.view(-1, prompt_image_embeds.shape[-1])
+
+            visual_idx = image_idx
+
+            prompt_image_embeds = prompt_image_embeds.to(dtype=text_embeds.dtype)
+
+            text_embeds[image_idx == 1] = torch.zeros_like(text_embeds[image_idx == 1]) + prompt_image_embeds.to(text_embeds.device)
+
+
+        elif video is not None:
+            if len(video.shape) == 5:
+                B, T, C, H, W = video.shape
+                N = 1
+                if self.use_clip:
+                    video = video.reshape(B*self.num_clips, T//self.num_clips, C, H, W)  # [16, 8, 3, 224, 224]
+            else:
+                B, N, T, C, H, W = video.shape
+
+            video = video.permute(0,2,1,3,4)      # 
+
+
+            prompt_video_embeds = self.encode_vision(video, instruction=time_instructions)  # [2, 96, 768]
+            if self.use_clip:
+                prompt_video_embeds = prompt_video_embeds.reshape(B,-1,prompt_video_embeds.shape[-1])     # [2,8*96,768]
+                batch_size, img_len, token_dim = prompt_video_embeds.shape
+                prompt_video_embeds = prompt_video_embeds.view(batch_size, img_len // self.token_merge_len, self.token_merge_len * token_dim)  # [B, 768//4, 4*768] = [2, 192, 3072]
+                prompt_video_embeds = self.merge_proj(prompt_video_embeds)   # [2, 192, 4096]
+                prompt_video_embeds = prompt_video_embeds.view(-1, prompt_video_embeds.shape[-1]) # [2*192, 4096]
+        
+            else:
+                prompt_video_embeds = self.project_up(prompt_video_embeds) # [2, 96, 4096]
+
+            prompt_video_embeds = prompt_video_embeds.view(-1, prompt_video_embeds.shape[-1]) 
+            visual_idx = video_idx
+            
+         
+            text_embeds[video_idx == 1] = torch.zeros_like(text_embeds[video_idx == 1]) + prompt_video_embeds.to(text_embeds.device).to(text_embeds.dtype)
+        
+        else:
+            logger.warn(f"don't get visual input, input_ids: {input_ids}")    
+        
+
+        for idx, text_embed in enumerate(text_embeds): 
+            if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.box_token].shape[0] != 0:
+                text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.box_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.box_token]) + torch.cat([self.box_token.squeeze(0)] * (text_embeds[idx][input_ids[idx] == self.tokenizer.box_token]).shape[0]).to(text_embeds.dtype)
+            if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.temp_token].shape[0] != 0:            
+                text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.temp_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.temp_token]) + self.temporal_token
+            if text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.track_token].shape[0] != 0:            
+                text_embeds[idx][input_ids[idx].flatten() == self.tokenizer.track_token] = torch.zeros_like(text_embeds[idx][input_ids[idx] == self.tokenizer.track_token]) + self.track_token
+             
+        if return_visual:
+            return text_embeds, visual, visual_idx
+        
+        return text_embeds
+
+ 
+    
+    def temporal_decode(self, temporal_embedding):
+        pred_sted = self.temporal_embed(temporal_embedding)
+        pred_actioness = self.action_embed(temporal_embedding)
+        return pred_sted, pred_actioness
+
+  
+    def box_loss2(self, src_boxes, target_boxes):
+        src_boxes = torch.cat(src_boxes)
+        target_boxes = torch.cat(target_boxes)
+
+        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
+        loss_bbox = self.masked_loss(loss_bbox, 0)
+        mask = (src_boxes[2:] >= src_boxes[:2]).all(-1)
+        src_boxes = src_boxes[mask]
+        target_boxes = target_boxes[mask]
+
+        return loss_bbox 
+
+    def box_loss(self, src_boxes, target_boxes):
+        src_boxes = torch.cat(src_boxes)
+        target_boxes = torch.cat(target_boxes)
+
+        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
+        loss_bbox = self.masked_loss(loss_bbox, 0)
+        mask = (src_boxes[:, 2:] >= src_boxes[ :, :2]).all(-1)
+        src_boxes = src_boxes[mask]
+        target_boxes = target_boxes[mask]
+
+        if src_boxes.shape[0] > 0:
+            loss_giou = 1 - torch.diag(generalized_box_iou(
+                src_boxes,
+                target_boxes))
+            loss_giou = self.masked_loss(loss_giou, 0)
+        else:
+            loss_giou = torch.tensor(2, dtype=src_boxes.dtype)
+        iou, union = box_iou(src_boxes, target_boxes)
+        
+        return loss_bbox * 2 + loss_giou / 5
+    
+    def find_boundaries_torch(self, mask):
+
+        from skimage.segmentation import find_boundaries
+        mask_np = mask.to(torch.bool).numpy()
+        boundaries = find_boundaries(mask_np, mode='outer')
+        boundary_points = np.argwhere(boundaries)
+        if boundary_points.size == 0:
+            return torch.tensor([-1, -1, -1, -1], dtype = torch.bfloat16)
+        h0, w0 = boundary_points.min(axis=0)
+        h1, w1 = boundary_points.max(axis=0)
+        return torch.tensor([w0 / mask.shape[1], h0 / mask.shape[0],  w1 / mask.shape[1], h1 / mask.shape[0]], dtype = torch.bfloat16)
+
+
+    def sam_loss(self, sam_outputs, gt_masks):
+        bound1 = self.find_boundaries_torch(gt_masks[:,:,:1].squeeze(2).cpu())
+        bound2 = self.find_boundaries_torch(sam_outputs[:,:,:1].squeeze(2).cpu()) 
+
+        lossl1 = F.l1_loss(bound1, bound2, reduction='none')
+        lossl1 = self.masked_loss(lossl1, 0)
+
+        loss_iou = self.iou_loss(sam_outputs, gt_masks)
+        loss_dice = self.dice_loss(sam_outputs, gt_masks)
+
+        # print(f'mask loss:{loss_iou,  loss_dice}')
+        return loss_iou + loss_dice + lossl1
+    
+    def masked_loss(self, loss, n):
+        mask = torch.ones_like(loss)
+        # mask[-n:] = 1e-10
+        loss = (loss*mask).sum()/(mask.sum())
+        return loss
+
+    def encode_vision(
+        self,
+        image,
+        instruction
+    ):
+        device = image.device
+        B = image.shape[0]
+        T = image.shape[2]
+        use_image = True if T == 1 else False
+        image_embeds = self.vision_encoder(image, use_image=use_image)
+        C = image_embeds.shape[-1]
+        image_embeds = image_embeds.reshape(B, -1, C)
+        image_embeds = self.vision_layernorm(image_embeds).to(device)  # [B, T*L, C]
+        
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(device)
+        if self.extra_num_query_token > 0:
+            query_tokens = torch.cat([self.query_tokens, self.extra_query_tokens], dim=1)
+        query_tokens = query_tokens.expand(image_embeds.shape[0], -1, -1)
+        if instruction is not None:
+            text_Qformer = self.qformer_tokenizer(
+                instruction,
+                padding='longest',
+                truncation=True,
+                max_length=512,
+                return_tensors="pt",
+            ).to(image_embeds.device)
+            query_atts = torch.ones(query_tokens.size()[:-1], dtype=torch.long).to(image_embeds.device)
+            Qformer_atts = torch.cat([query_atts, text_Qformer.attention_mask], dim=1)
+            query_output = self.qformer.bert(
+                text_Qformer.input_ids,
+                attention_mask=Qformer_atts,
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+        else:
+            query_output = self.qformer.bert(
+                query_embeds=query_tokens,
+                encoder_hidden_states=image_embeds,
+                encoder_attention_mask=image_atts,
+                return_dict=True,
+            )
+        
+        return query_output.last_hidden_state[:, :query_tokens.size(1), :]
+
+    def generate_caption(
+        self,
+        input_ids,
+        attention_mask,
+        image_idx = None,
+        video_idx = None,
+        image: Optional[torch.Tensor] = None,
+        video: Optional[torch.Tensor] = None,
+        num_beams=1,
+        max_new_tokens=200,
+        do_sample=True,
+        top_p=0.9,
+        top_k=None,
+        temperature=1.0,
+        length_penalty=1,
+        repetition_penalty=1.0,
+    ):
+        text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx)
+        outputs = self.lm.generate(
+            inputs_embeds=text_embeds,
+            attention_mask=attention_mask,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            min_length=1,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            length_penalty=length_penalty,
+            repetition_penalty=repetition_penalty,
+        )
+
+        return outputs
+
+    def generate_caption_bbox(
+        self,
+        input_ids,
+        attention_mask,
+        labels,
+        image_idx = None,
+        video_idx = None,
+        image: Optional[torch.Tensor] = None,
+        video: Optional[torch.Tensor] = None,
+        num_beams=1,
+        max_new_tokens=200,
+        do_sample=True,
+        top_p=0.9,
+        top_k=None,
+        temperature=0.9,
+        length_penalty=1,
+        repetition_penalty=1.0,
+    ):
+        text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx)
+        outputs = self.lm.generate(
+            inputs_embeds=text_embeds,
+            attention_mask=attention_mask,
+            num_beams=num_beams,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            min_length=1,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            length_penalty=length_penalty,
+            repetition_penalty=repetition_penalty,
+        )
+        decoded_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # torch.save({'text':decoded_text, 'output':{outputs}}, 'tmp.pth')
+        # print(decoded_text)
+        return outputs
+    
+    def generate_temporal(self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        image: Optional[torch.Tensor] = None,
+        video: Optional[torch.Tensor] = None,
+        instruction = None,
+        video_idx = None,
+        image_idx = None,
+        boxes = None,
+        text_input = None,
+        video_info = None, 
+        temporal_labels = None):
+
+        if text_input is not None:
+            time_instructions = self.get_clip_time_instruct(text_input)
+        else:
+            time_instructions = None
+        text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, return_visual=False,
+                                        video_idx=video_idx, image_idx=image_idx,  instruction = instruction, 
+                                        boxes = boxes, time_instructions = time_instructions)
+        
+        # TODO
+        outputs = self.lm(
+            inputs_embeds=text_embeds,
+            attention_mask=attention_mask,
+            labels=labels,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+
+        if temporal_labels is not None:
+            start_sec = temporal_labels["start_sec"]
+            end_sec = temporal_labels["end_sec"]
+            fps = video_info['fps']
+            frame_indices = video_info['frame_indices']
+
+            last_hidden_states = outputs.hidden_states[-1]                                   # [2,1024, 4096]
+            last_hidden_states = last_hidden_states.view(-1, last_hidden_states.size(-1))    # [2048, 4096]
+            loc_positions = (input_ids.flatten()==self.tokenizer.temp_place_ids).nonzero().flatten()   #
+            selected_hidden_states = last_hidden_states[loc_positions]
+            selected_hidden_states = selected_hidden_states.view(input_ids.shape[0], -1 ,selected_hidden_states.shape[-1]) # [2, 64, 4096]
+
+            # just for debug
+            
+            # vis_embed = vis_embed[:,:64,:]
+
+            pred_sted, pred_actionness = self.temporal_decode(selected_hidden_states) # [2,64,2]  [2,64,1]
+
+            pred_sted = self.postprocess(pred_sted, frame_indices)
+            pred_sec_s = pred_sted[0][0] / fps[0][0].item()
+            pred_sec_e = pred_sted[0][1] / fps[0][0].item()
+
+            output_file = "predictions2.jsonl"
+            prediction = {"pred_sec_s": round(pred_sec_s, 1), "pred_sec_e": round(pred_sec_e, 1), "start_sec":float(start_sec[0]), "end_sec": float(end_sec[0])}
+
+            with open(output_file, 'a') as f:
+                json.dump(prediction, f)
+                f.write('\n')
+
+            return outputs
+
+    def generate_seg(self, input_ids, attention_mask, labels, image, image_idx, video, video_idx, input_boxes, size_hw, sam_images):
+        device = input_ids.device
+        prompt = input_ids
+        l_prompt = len(input_ids)
+        temperature = 1e-5
+        max_new_tokens = 20
+        guide_w = 5
+        stop_str = '</s>'
+        bbox = []
+        output_ids = list(input_ids[0])
+        text_embeds = self.pad_text_embeds(input_ids=input_ids, image=image, video=video, image_idx=image_idx, video_idx=video_idx, return_visual=False, 
+                                        instruction = None, output_boxes=None, input_boxes=input_boxes)
+        for i in range(max_new_tokens):
+            if i == 0:
+                outputs = self.lm(
+                        inputs_embeds=text_embeds,
+                        attention_mask=attention_mask,
+                        output_hidden_states=True,
+                        return_dict=True,
+                    )
+                logits = outputs.logits
+                past_key_values = outputs.past_key_values
+            else:
+                attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
+                last_text_embeds = self.lm.get_input_embeddings()(torch.tensor(output_ids[-1], device=device).long()).detach().unsqueeze(0)
+                last_text_embeds = last_text_embeds.unsqueeze(0)
+                
+                out = self.lm(
+                    input_ids=None,
+                    use_cache=True,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                    inputs_embeds=last_text_embeds,
+                    past_key_values=past_key_values,
+                )
+                logits = out.logits
+                past_key_values = out.past_key_values
+            if logits is not None:
+                last_token_logits = logits[0][-1]
+                if temperature < 1e-4:
+                    token = int(torch.argmax(last_token_logits))
+                else:
+                    probs = torch.softmax(last_token_logits / temperature, dim=-1)
+                    token = int(torch.multinomial(probs, num_samples=1))
+                output_ids.append(token)
+            ret = self.tokenizer.decode(token)
+            if ret == '<box_begin>':
+                attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
+                bbox_embeds = self.box_token.bfloat16()
+                out = self.lm(
+                    inputs_embeds=bbox_embeds,
+                    use_cache=True,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                    past_key_values=past_key_values
+                )
+                last_hidden_states = out.hidden_states[-1]
+                selected_hidden_states = last_hidden_states[0][0]
+                bbox.append(self.loc_decoder(selected_hidden_states))
+                last_token_logits = logits[0][-1]
+                if temperature < 1e-4:
+                    token = int(torch.argmax(last_token_logits))
+                else:
+                    probs = torch.softmax(last_token_logits / temperature, dim=-1)
+                    token = int(torch.multinomial(probs, num_samples=1))
+            if ret == '<track_begin>':
+                attention_mask = torch.ones(1, past_key_values[0][0].shape[-2] + 1, device=device)
+                tracking_embeds = self.track_token
+                out = self.lm(
+                    inputs_embeds=tracking_embeds,
+                    use_cache=True,
+                    attention_mask=attention_mask,
+                    output_hidden_states=True,
+                    past_key_values=past_key_values
+                )
+                last_hidden_states = out.hidden_states[-1]
+                selected_hidden_states = last_hidden_states[0][0].to(dtype = torch.bfloat16)
+                
+                embed_sam_boxes = self.track_embed(selected_hidden_states).reshape(1, 3, 256)
+                
+                inference_state = self.sam.init_state_images(sam_images, size_hw[0][0], size_hw[0][1])
+                gt_embeds = self.sam.get_prompt_embeding(inference_state, None, None, False, input_boxes[0].cuda(), device = text_embeds.device) 
+                ann_frame_idx = 0
+                ann_obj_id = 0
+                box = np.array([0, 0, 0, 0], dtype=np.float32)
+                _, out_obj_ids, out_mask_logits = self.sam.add_new_box_embeding(
+                    inference_state=inference_state,
+                    frame_idx=ann_frame_idx,
+                    obj_id=ann_obj_id,
+                    box=box,
+                    box_embeding=embed_sam_boxes,
+                )
+                video_segments = {}  # video_segments contains the per-frame segmentation results
+                for out_frame_idx, out_obj_ids, out_mask_logits in self.sam.propagate_in_video(inference_state):
+                    video_segments[out_frame_idx] = {
+                        out_obj_id: (out_mask_logits[i] > 0.0)
+                        for i, out_obj_id in enumerate(out_obj_ids)
+                    }
+                video_segments = [video_segments[tt][0] for tt in video_segments]
+                # bbox = model.find_boundaries_torch(video_segments[0].squeeze(0).cpu())
+                # return ret, [], video_segments
+
+            if (ret == '</s>'):
+                break
+        ret = self.tokenizer.decode(output_ids)
+        del past_key_values
+        return ret, bbox, video_segments
+    
+    def generate_answer(self, tokenizer, instruction, msg, user_prompt, media_type="video",video_tensor=None, image_tensor=None, answer_prompt=None, chat_history=[],return_history=False, debug=False, generation_config={}):
+        input_ids, attention_masks, labels = [], [], []
+
+        conversation = ""
+        if instruction:
+            conversation += instruction
+        conversation += (
+                    "[INST]" + " "
+                )
+
+        if media_type == 'image':
+            conversation +=( "<Image>" + IMG_TOKEN + "</Image>")
+        else:
+            conversation += ("<Video>" + VID_TOKEN + "</Video>")
+
+        conversation += ( msg.rstrip() + "[/INST]")
+
+        for q,a in chat_history:
+            conversation += (" [INST] " + q + " [/INST]")
+            conversation += (a + "</s>")
+            
+        conversation += (" [INST] " + user_prompt + " [/INST]")
+        conversation += ("")
+        if answer_prompt:
+            conversation += ("Best Option: (")
+        total_len = 0
+        indexs = []
+        if debug:
+            print(conversation)
+
+        tokenized = tokenizer.build_input_ids([conversation], 
+                                          max_length=1024, 
+                                          add_special_tokens=True, 
+                                          truncation=False, 
+                                          padding=False, 
+                                          return_tensors='pt', 
+                                          image=image_tensor,
+                                          video=video_tensor,
+                                          require_video=True)
+        if video_tensor is not None:
+            generation_output = self.generate_caption(
+                    tokenized['input_ids'].unsqueeze(0).to(self.device), 
+                    tokenized['attention_mask'].unsqueeze(0).to(self.device), 
+                    video_idx = tokenized['video_index'].unsqueeze(0),
+                    video = video_tensor.unsqueeze(0).to(self.device,dtype=torch.bfloat16), 
+                    do_sample=False
+                    )
+        elif image_tensor is not None:
+            generation_output = self.generate_caption(
+                    tokenized['input_ids'].unsqueeze(0).to(self.device), 
+                    tokenized['attention_mask'].unsqueeze(0).to(self.device), 
+                    image_idx = tokenized['image_index'].unsqueeze(0),
+                    image = image_tensor.unsqueeze(0).to(self.device,dtype=torch.bfloat16), 
+                    do_sample=False
+                    )
+        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
+        if debug:
+            print(response)
+        return response, chat_history

modeling_vit.py

@@ -0,0 +1,487 @@
+import logging
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from functools import partial
+
+from timm.models.layers import drop_path, to_2tuple, trunc_normal_
+
+logger = logging.getLogger(__name__)
+
+
+def _cfg(url='', **kwargs):
+    return {
+        'url': url,
+        'num_classes': 400, 'input_size': (3, 224, 224), 'pool_size': None,
+        'crop_pct': .9, 'interpolation': 'bicubic',
+        'mean': (0.5, 0.5, 0.5), 'std': (0.5, 0.5, 0.5),
+        **kwargs
+    }
+
+class MLP(nn.Module):
+    """Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, dropout=0):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(
+            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])
+        )
+        self.dropout = dropout
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+            if self.dropout and i < self.num_layers:
+                x = self.dropout(x)
+        return x
+
+class PostProcess(nn.Module):
+    """ This module converts the model's output into the format expected by the coco api"""
+    
+    @torch.no_grad()
+    def forward(self, out_sted, frames_id):
+        """Perform the computation for inference evaluation
+        """
+        # import pdb; pdb.set_trace()
+
+        b, t, _ = out_sted.shape
+        device = out_sted.device
+        temp_prob_map = torch.zeros(b,t,t).to(device)
+        inf = -1e32
+        for i_b in range(len(frames_id)): 
+            duration = len(frames_id[0])
+            sted_prob = (torch.ones(t, t) * inf).tril(0).to(device)
+            sted_prob[duration:,:] = inf
+            sted_prob[:,duration:] = inf
+            temp_prob_map[i_b,:,:] = sted_prob
+        
+        temp_prob_map += F.log_softmax(out_sted[:, :, 0], dim=1).unsqueeze(2) + \
+                F.log_softmax(out_sted[:, :, 1], dim=1).unsqueeze(1)
+        
+        pred_steds = []
+        for i_b in range(b):
+            prob_map = temp_prob_map[i_b]  # [T * T]
+            frame_id_seq = frames_id[i_b]
+            prob_seq = prob_map.flatten(0)  
+            max_tstamp = prob_seq.max(dim=0)[1].item()
+            start_idx = max_tstamp // t
+            end_idx = max_tstamp % t
+            pred_sted = [frame_id_seq[start_idx], frame_id_seq[end_idx]+1]
+            pred_steds.append(pred_sted)
+    
+        return pred_steds
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+    
+    def extra_repr(self) -> str:
+        return 'p={}'.format(self.drop_prob)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+            self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.,
+            proj_drop=0., attn_head_dim=None):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        if attn_head_dim is not None:
+            head_dim = attn_head_dim
+        all_head_dim = head_dim * self.num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, all_head_dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(all_head_dim))
+            self.v_bias = nn.Parameter(torch.zeros(all_head_dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(all_head_dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        # qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q @ k.transpose(-2, -1))
+        
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., init_values=None, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
+                 attn_head_dim=None):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+            attn_drop=attn_drop, proj_drop=drop, attn_head_dim=attn_head_dim)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)),requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x):
+        if self.gamma_1 is None:
+            x = x + self.drop_path(self.attn(self.norm1(x)))
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x)))
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, num_frames=16, tubelet_size=2):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.tubelet_size = int(tubelet_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) * (num_frames // self.tubelet_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = nn.Conv3d(
+            in_channels=in_chans, out_channels=embed_dim, 
+            kernel_size=(self.tubelet_size, patch_size[0], patch_size[1]), 
+            stride=(self.tubelet_size, patch_size[0], patch_size[1])
+        )
+        logger.info(f'Num of patches: {num_patches}')
+
+    def forward(self, x, **kwargs):
+        B, C, T, H, W = x.shape
+        # FIXME look at relaxing size constraints
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+    
+# sin-cos position encoding
+# https://github.com/jadore801120/attention-is-all-you-need-pytorch/blob/master/transformer/Models.py#L31
+def get_sinusoid_encoding_table(n_position, d_hid, ckpt_num_frame=-1, cur_frame=12): 
+    ''' Sinusoid position encoding table ''' 
+    # TODO: make it with torch instead of numpy 
+    def get_position_angle_vec(position): 
+        return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] 
+    
+    if ckpt_num_frame != -1 and ckpt_num_frame != cur_frame:
+        logger.info(f"Interpolate position embedding")
+        logger.info(f"Testing frame: {cur_frame}")
+        logger.info(f"Checkpoint frame: {ckpt_num_frame}")
+
+        T = ckpt_num_frame # checkpoint frame
+        new_T = cur_frame # testing frame
+        n_position = n_position // new_T * T # generate checkpoint position embedding
+        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) 
+        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 
+        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 
+        sinusoid_table = torch.tensor(sinusoid_table, dtype=torch.float, requires_grad=False).unsqueeze(0)
+        # interpolate
+        P = int((n_position // T) ** 0.5)
+        C = d_hid
+        sinusoid_table = sinusoid_table.reshape(-1, T, P, P, C)
+        sinusoid_table = sinusoid_table.permute(0, 2, 3, 4, 1).reshape(-1, C, T)  # BHW, C, T
+        sinusoid_table = torch.nn.functional.interpolate(sinusoid_table, size=new_T, mode='linear')
+        sinusoid_table = sinusoid_table.reshape(1, P, P, C, new_T).permute(0, 4, 1, 2, 3) # B, T, H, W, C
+        sinusoid_table = sinusoid_table.flatten(1, 3)
+        return sinusoid_table
+    else:
+        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) 
+        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 
+        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 
+        return torch.tensor(sinusoid_table, dtype=torch.float, requires_grad=False).unsqueeze(0) 
+    
+
+def get_sinusoid_encoding_table2(n_position=784, d_hid=1024, cur_frame=8, ckpt_num_frame=4, pre_n_position=784): 
+    ''' Sinusoid position encoding table ''' 
+    # TODO: make it with torch instead of numpy 
+    def get_position_angle_vec(position): 
+        return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)] 
+    
+    # generate checkpoint position embedding
+    sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(pre_n_position)]) 
+    sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i 
+    sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1 
+    sinusoid_table = torch.tensor(sinusoid_table, dtype=torch.float, requires_grad=False).unsqueeze(0)
+    
+    print(f"n_position: {n_position}")
+    print(f"pre_n_position: {pre_n_position}")
+    
+    if n_position != pre_n_position:
+        T = ckpt_num_frame # checkpoint frame
+        P = 14 # checkpoint size
+        C = d_hid
+        new_P = int((n_position // cur_frame) ** 0.5) # testing size
+        print(f'Pretraining uses 14x14, but current version is {new_P}x{new_P}')
+        print(f'Interpolate the position embedding')
+        sinusoid_table = sinusoid_table.reshape(-1, T, P, P, C)
+        sinusoid_table = sinusoid_table.reshape(-1, P, P, C).permute(0, 3, 1, 2)
+        sinusoid_table = torch.nn.functional.interpolate(
+            sinusoid_table, size=(new_P, new_P), mode='bicubic', align_corners=False)
+        # BT, C, H, W -> BT, H, W, C ->  B, T, H, W, C
+        sinusoid_table = sinusoid_table.permute(0, 2, 3, 1).reshape(-1, T, new_P, new_P, C)
+        sinusoid_table = sinusoid_table.flatten(1, 3)  # B, THW, C
+    
+    if cur_frame != ckpt_num_frame:
+        print(f'Pretraining uses 4 frames, but current frame is {cur_frame}')
+        print(f'Interpolate the position embedding')
+        T = ckpt_num_frame # checkpoint frame
+        new_T = cur_frame # testing frame
+        # interpolate
+        P = int((n_position // cur_frame) ** 0.5) # testing size
+        C = d_hid
+        sinusoid_table = sinusoid_table.reshape(-1, T, P, P, C)
+        sinusoid_table = sinusoid_table.permute(0, 2, 3, 4, 1).reshape(-1, C, T)  # BHW, C, T
+        sinusoid_table = torch.nn.functional.interpolate(sinusoid_table, size=new_T, mode='linear')
+        sinusoid_table = sinusoid_table.reshape(1, P, P, C, new_T).permute(0, 4, 1, 2, 3) # B, T, H, W, C
+        sinusoid_table = sinusoid_table.flatten(1, 3)  # B, THW, C
+        
+    return sinusoid_table
+
+
+class PretrainVisionTransformerEncoder(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, init_values=None, num_frames=8, tubelet_size=1,
+                 use_learnable_pos_emb=False,
+                 use_checkpoint=False, checkpoint_num=0, 
+                 ckpt_num_frame=-1, with_ln=True, return_index=-1
+                 ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+        self.patch_embed = PatchEmbed(
+            img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim, 
+            num_frames=num_frames, tubelet_size=tubelet_size
+        )
+        num_patches = self.patch_embed.num_patches
+        self.depth = depth + return_index + 1
+        self.use_checkpoint = use_checkpoint
+        self.checkpoint_num = checkpoint_num
+        logger.info(f"Use checkpoint: {use_checkpoint}")
+        logger.info(f"Checkpoint number: {checkpoint_num}")
+        logger.info(f"Real runing depth: {self.depth}")
+
+        # TODO: Add the cls token
+        if use_learnable_pos_emb:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+            self.img_pos_embed = nn.Parameter(torch.zeros(1, num_patches//(num_frames//tubelet_size) + 1, embed_dim))
+        else:
+            # sine-cosine positional embeddings 
+            if img_size != 224:
+                self.pos_embed = get_sinusoid_encoding_table2(num_patches, embed_dim, ckpt_num_frame=ckpt_num_frame, cur_frame=num_frames//tubelet_size)
+                self.img_pos_embed = get_sinusoid_encoding_table2(num_patches//(num_frames//tubelet_size), embed_dim, cur_frame=1, ckpt_num_frame=1, pre_n_position=14*14)
+            else:
+                self.pos_embed = get_sinusoid_encoding_table(num_patches, embed_dim, ckpt_num_frame=ckpt_num_frame, cur_frame=num_frames//tubelet_size)
+                self.img_pos_embed = get_sinusoid_encoding_table(num_patches//(num_frames//tubelet_size), embed_dim)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                init_values=init_values)
+            for i in range(self.depth)])
+        
+        if with_ln:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = nn.Identity()
+
+        if use_learnable_pos_emb:
+            trunc_normal_(self.pos_embed, std=.02)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def forward_features(self, x, use_image=False):
+        x = self.patch_embed(x)
+        
+        if use_image:
+            x = x + self.img_pos_embed.type_as(x).to(x.device).clone().detach()
+        else:
+            x = x + self.pos_embed.type_as(x).to(x.device).clone().detach()
+
+        B, _, C = x.shape
+        x_vis = x
+
+        for idx, blk in enumerate(self.blocks):
+            if self.use_checkpoint and idx < self.checkpoint_num:
+                x_vis = checkpoint.checkpoint(blk, x_vis)
+            else:
+                x_vis = blk(x_vis)
+
+        # with ln ot not
+        x_vis = self.norm(x_vis)
+        return x_vis
+
+    def forward(self, x, use_image=False):
+        x_vis = self.forward_features(x, use_image)
+        return x_vis
+
+
+class PretrainVisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+    def __init__(self,
+                 img_size=224, 
+                 patch_size=16, 
+                 encoder_in_chans=3, 
+                 encoder_embed_dim=768, 
+                 encoder_depth=12,
+                 encoder_num_heads=12, 
+                 mlp_ratio=4., 
+                 qkv_bias=True, 
+                 qk_scale=None, 
+                 drop_rate=0., 
+                 attn_drop_rate=0.,
+                 drop_path_rate=0., 
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6), 
+                 init_values=0.,
+                 use_learnable_pos_emb=False,
+                 num_frames=8,
+                 tubelet_size=1,
+                 use_checkpoint=False,
+                 checkpoint_num=0,
+                 ckpt_num_frame=4, # the pretrained model uses 4 frames
+                 return_index=-1,
+                 with_ln=False
+                ):
+        super().__init__()
+
+        self.encoder = PretrainVisionTransformerEncoder(
+            img_size=img_size, 
+            patch_size=patch_size, 
+            in_chans=encoder_in_chans, 
+            embed_dim=encoder_embed_dim, 
+            depth=encoder_depth,
+            num_heads=encoder_num_heads, 
+            mlp_ratio=mlp_ratio, 
+            qkv_bias=qkv_bias, 
+            qk_scale=qk_scale, 
+            drop_rate=drop_rate, 
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate, 
+            norm_layer=norm_layer, 
+            init_values=init_values,
+            num_frames=num_frames,
+            tubelet_size=tubelet_size,
+            use_learnable_pos_emb=use_learnable_pos_emb,
+            use_checkpoint=use_checkpoint,
+            checkpoint_num=checkpoint_num,
+            ckpt_num_frame=ckpt_num_frame,
+            with_ln=with_ln,
+            return_index=return_index
+        )
+        logger.info(f'With LN: {with_ln}')
+        logger.info(f'Total {encoder_depth} layer')
+        logger.info(f'Return {encoder_depth+return_index+1}-th layer')
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token', 'clip_pos_embed'}
+
+    def forward(self, x, use_image=False):
+        T = x.shape[2]
+        x_vis = self.encoder(x, use_image) # [B, N_vis, C_e]
+        B, TL, C = x_vis.shape
+        x_vis = x_vis.view(B, T, TL // T, C)
+
+        return x_vis
+
+
+def build_vit(config):
+    model = PretrainVisionTransformer(
+        img_size=config.vision_encoder.img_size, 
+        patch_size=config.vision_encoder.patch_size, 
+        encoder_embed_dim=config.vision_encoder.encoder_embed_dim, 
+        encoder_depth=config.vision_encoder.encoder_depth,
+        encoder_num_heads=config.vision_encoder.encoder_num_heads, 
+        drop_path_rate=config.vision_encoder.drop_path_rate, 
+        num_frames=config.vision_encoder.num_frames,
+        tubelet_size=config.vision_encoder.tubelet_size,
+        use_checkpoint=config.vision_encoder.use_checkpoint,
+        checkpoint_num=config.vision_encoder.checkpoint_num,
+        return_index=config.vision_encoder.get('return_index', -1),
+        with_ln=config.vision_encoder.get('with_ln', False),
+    )
+    model.default_cfg = _cfg()
+    if config.vision_encoder.pretrained:
+        logger.info(f"Loading pretrained weights from {config.vision_encoder.pretrained}")
+        state_dict = torch.load(config.vision_encoder.pretrained, map_location='cpu')
+        model.load_state_dict(state_dict, strict=False)
+    else:
+        logger.info("No pretrained weights!!!")
+    return model
+ 

special_tokens_map.json

@@ -0,0 +1,24 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<unk>",
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

third_party/__init__.py

@@ -0,0 +1,2 @@
+import logging
+logger = logging.getLogger(__name__)

third_party/cgdetr/cg_detr/attention.py

@@ -0,0 +1,394 @@
+# ------------------------------------------------------------------------
+# DAB-DETR
+# Copyright (c) 2022 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+# Modified from codes in torch.nn
+# ------------------------------------------------------------------------
+
+"""
+MultiheadAttention that support query, key, and value to have different dimensions.
+Query, key, and value projections are removed.
+Mostly copy-paste from https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/activation.py#L873
+and https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4837
+"""
+
+import copy
+from typing import Optional, List
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+import warnings
+from typing import Tuple, Optional
+
+import torch
+from torch import Tensor
+from torch.nn.modules.linear import Linear
+from torch.nn.init import xavier_uniform_
+from torch.nn.init import constant_
+from torch.nn.init import xavier_normal_
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from torch.nn import functional as F
+
+import warnings
+import math
+
+from torch._C import _infer_size, _add_docstr
+from torch.nn import _reduction as _Reduction
+from torch.nn.modules import utils
+from torch.nn.modules.utils import _single, _pair, _triple, _list_with_default
+from torch.nn import grad
+from torch import _VF
+from torch._jit_internal import boolean_dispatch, List, Optional, _overload, Tuple
+try:
+    from torch.overrides import has_torch_function, handle_torch_function
+except:
+    from torch._overrides import has_torch_function, handle_torch_function
+Tensor = torch.Tensor
+
+from torch.nn.functional import linear, pad, softmax, dropout
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces.
+    See reference: Attention Is All You Need
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
+    Args:
+        embed_dim: total dimension of the model.
+        num_heads: parallel attention heads.
+        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
+        bias: add bias as module parameter. Default: True.
+        add_bias_kv: add bias to the key and value sequences at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        kdim: total number of features in key. Default: None.
+        vdim: total number of features in value. Default: None.
+        Note: if kdim and vdim are None, they will be set to embed_dim such that
+        query, key, and value have the same number of features.
+    Examples::
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+    """
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
+        super(MultiheadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        vdim = vdim if vdim is not None else embed_dim
+        self.out_proj = Linear(vdim , vdim)
+
+        self.in_proj_bias = None
+        self.in_proj_weight = None
+        self.bias_k = self.bias_v = None
+        self.q_proj_weight = None
+        self.k_proj_weight = None
+        self.v_proj_weight = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.out_proj.bias, 0.)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if '_qkv_same_embed_dim' not in state:
+            state['_qkv_same_embed_dim'] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(self, query, key, value, key_padding_mask=None,
+                need_weights=True, attn_mask=None):
+        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
+        r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. When given a binary mask and a value is True,
+            the corresponding value on the attention layer will be ignored. When given
+            a byte mask and a value is non-zero, the corresponding value on the attention
+            layer will be ignored
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+    Shape:
+        - Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the position
+          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*\text{num_heads}, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+        """
+        if not self._qkv_same_embed_dim:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight, out_dim=self.vdim)
+        else:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, out_dim=self.vdim)
+
+
+def multi_head_attention_forward(query: Tensor,
+                                 key: Tensor,
+                                 value: Tensor,
+                                 embed_dim_to_check: int,
+                                 num_heads: int,
+                                 in_proj_weight: Tensor,
+                                 in_proj_bias: Tensor,
+                                 bias_k: Optional[Tensor],
+                                 bias_v: Optional[Tensor],
+                                 add_zero_attn: bool,
+                                 dropout_p: float,
+                                 out_proj_weight: Tensor,
+                                 out_proj_bias: Tensor,
+                                 training: bool = True,
+                                 key_padding_mask: Optional[Tensor] = None,
+                                 need_weights: bool = True,
+                                 attn_mask: Optional[Tensor] = None,
+                                 use_separate_proj_weight: bool = False,
+                                 q_proj_weight: Optional[Tensor] = None,
+                                 k_proj_weight: Optional[Tensor] = None,
+                                 v_proj_weight: Optional[Tensor] = None,
+                                 static_k: Optional[Tensor] = None,
+                                 static_v: Optional[Tensor] = None,
+                                 out_dim: Optional[Tensor] = None
+                                 ) -> Tuple[Tensor, Optional[Tensor]]:
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+    Shape:
+        Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
+          will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+    """
+    if not torch.jit.is_scripting():
+        tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,
+                    out_proj_weight, out_proj_bias)
+        if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
+            return handle_torch_function(
+                multi_head_attention_forward, tens_ops, query, key, value,
+                embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,
+                bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,
+                out_proj_bias, training=training, key_padding_mask=key_padding_mask,
+                need_weights=need_weights, attn_mask=attn_mask,
+                use_separate_proj_weight=use_separate_proj_weight,
+                q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,
+                v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)
+    tgt_len, bsz, embed_dim = query.size()
+    assert embed_dim == embed_dim_to_check
+    # allow MHA to have different sizes for the feature dimension
+    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
+
+    head_dim = embed_dim // num_heads
+    v_head_dim = out_dim // num_heads
+    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+    scaling = float(head_dim) ** -0.5
+
+    q = query * scaling
+    k = key
+    v = value
+
+    if attn_mask is not None:
+        assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
+            attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
+            'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
+        if attn_mask.dtype == torch.uint8:
+            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+            attn_mask = attn_mask.to(torch.bool)
+
+        if attn_mask.dim() == 2:
+            attn_mask = attn_mask.unsqueeze(0)
+            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 2D attn_mask is not correct.')
+        elif attn_mask.dim() == 3:
+            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 3D attn_mask is not correct.')
+        else:
+            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
+        # attn_mask's dim is 3 now.
+
+    # convert ByteTensor key_padding_mask to bool
+    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+        warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+        key_padding_mask = key_padding_mask.to(torch.bool)
+
+    if bias_k is not None and bias_v is not None:
+        if static_k is None and static_v is None:
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if k is not None:
+        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    if v is not None:
+        v = v.contiguous().view(-1, bsz * num_heads, v_head_dim).transpose(0, 1)
+
+    if static_k is not None:
+        assert static_k.size(0) == bsz * num_heads
+        assert static_k.size(2) == head_dim
+        k = static_k
+
+    if static_v is not None:
+        assert static_v.size(0) == bsz * num_heads
+        assert static_v.size(2) == v_head_dim
+        v = static_v
+
+    src_len = k.size(1)
+
+    if key_padding_mask is not None:
+        assert key_padding_mask.size(0) == bsz
+        assert key_padding_mask.size(1) == src_len
+
+    if add_zero_attn:
+        src_len += 1
+        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
+        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_output_weights.masked_fill_(attn_mask, float('-inf'))
+        else:
+            attn_output_weights += attn_mask
+
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2),
+            float('-inf'),
+        )
+        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
+
+    # attn_output_weights = softmax(
+    #     attn_output_weights, dim=-1)
+    attn_output_weights = softmax(
+            attn_output_weights - attn_output_weights.max(dim=-1, keepdim=True)[0], dim=-1)
+    attn_output_weights = dropout(attn_output_weights, p=dropout_p, training=training)
+
+    attn_output = torch.bmm(attn_output_weights, v)
+    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, v_head_dim]
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, out_dim)
+    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        return attn_output, attn_output_weights.sum(dim=1) / num_heads
+    else:
+        return attn_output, None

third_party/cgdetr/cg_detr/config.py

@@ -0,0 +1,261 @@
+import os
+import time
+import torch
+import argparse
+
+from third_party.cgdetr.utils.basic_utils import mkdirp, load_json, save_json, make_zipfile, dict_to_markdown
+import shutil
+
+class BaseOptions(object):
+    saved_option_filename = "opt.json"
+    ckpt_filename = "model.ckpt"
+    tensorboard_log_dir = "tensorboard_log"
+    train_log_filename = "train.log.txt"
+    eval_log_filename = "eval.log.txt"
+
+    def __init__(self):
+        self.parser = None
+        self.initialized = False
+        self.opt = None
+
+    def initialize(self):
+        self.initialized = True
+        parser = argparse.ArgumentParser()
+        # parser.add_argument("--dset_name", type=str, choices=["hl", 'charadesSTA', ])
+        # parser.add_argument("--dset_domain", type=str, 
+        #                     help="Domain to train for tvsum dataset. (Only used for tvsum and youtube-hl)")
+        
+        parser.add_argument("--eval_split_name", type=str, default="val",
+                            help="should match keys in video_duration_idx_path, must set for VCMR")
+        parser.add_argument("--debug", action="store_true",
+                            help="debug (fast) mode, break all loops, do not load all data into memory.")
+        parser.add_argument("--data_ratio", type=float, default=1.0,
+                            help="how many training and eval data to use. 1.0: use all, 0.1: use 10%."
+                                 "Use small portion for debug purposes. Note this is different from --debug, "
+                                 "which works by breaking the loops, typically they are not used together.")
+        parser.add_argument("--results_root", type=str, default="results")
+        parser.add_argument("--exp_id", type=str, default=None, help="id of this run, required at training")
+        parser.add_argument("--seed", type=int, default=2018, help="random seed")
+        # parser.add_argument("--device", type=int, default=0, help="0 cuda, -1 cpu")
+        parser.add_argument("--num_workers", type=int, default=0,
+                            help="num subprocesses used to load the data, 0: use main process")
+        parser.add_argument("--no_pin_memory", action="store_true",
+                            help="Don't use pin_memory=True for dataloader. "
+                                 "ref: https://discuss.pytorch.org/t/should-we-set-non-blocking-to-true/38234/4")
+
+        # training config
+        # parser.add_argument("--lr", type=float, default=2e-4, help="learning rate")
+        # parser.add_argument("--lr_drop", type=int, default=800, help="drop learning rate to 1/10 every lr_drop epochs")
+        # parser.add_argument("--wd", type=float, default=1e-4, help="weight decay")
+        parser.add_argument("--n_epoch", type=int, default=200, help="number of epochs to run")
+        parser.add_argument("--max_es_cnt", type=int, default=200,
+                            help="number of epochs to early stop, use -1 to disable early stop")
+        # parser.add_argument("--bsz", type=int, default=32, help="mini-batch size")
+        # parser.add_argument("--eval_bsz", type=int, default=100,
+        #                     help="mini-batch size at inference, for query")
+        parser.add_argument("--eval_epoch", type=int, default=5,help="inference epoch")
+        parser.add_argument("--grad_clip", type=float, default=0.1, help="perform gradient clip, -1: disable")
+        parser.add_argument("--eval_untrained", action="store_true", help="Evaluate on un-trained model")
+        parser.add_argument("--resume", type=str, default=None,
+                            help="checkpoint path to resume or evaluate, without --resume_all this only load weights")
+        parser.add_argument("--resume_all", action="store_true",
+                            help="if --resume_all, load optimizer/scheduler/epoch as well")
+        parser.add_argument("--start_epoch", type=int, default=None,
+                            help="if None, will be set automatically when using --resume_all")
+
+        # Data config
+        parser.add_argument("--max_q_l", type=int, default=-1)
+        parser.add_argument("--max_v_l", type=int, default=-1)
+        parser.add_argument("--clip_length", type=float, default=2)
+        parser.add_argument("--max_windows", type=int, default=5)
+
+        parser.add_argument("--train_path", type=str, default=None)
+        parser.add_argument("--eval_path", type=str, default=None,
+                            help="Evaluating during training, for Dev set. If None, will only do training, ")
+        parser.add_argument("--no_norm_vfeat", action="store_true", help="Do not do normalize video feat")
+        parser.add_argument("--no_norm_tfeat", action="store_true", help="Do not do normalize text feat")
+        parser.add_argument("--v_feat_dirs", type=str, nargs="+",
+                            help="video feature dirs. If more than one, will concat their features. "
+                                 "Note that sub ctx features are also accepted here.")
+        parser.add_argument("--t_feat_dir", type=str, help="text/query feature dir")
+        # parser.add_argument("--a_feat_dir", type=str, help="audio feature dir")
+        parser.add_argument("--v_feat_dim", type=int, default=770, help="video feature dim")
+        parser.add_argument("--t_feat_dim", type=int, default=4096, help="text/query feature dim")
+        # parser.add_argument("--a_feat_dim", type=int, help="audio feature dim")
+        parser.add_argument("--ctx_mode", type=str, default="video_tef")
+
+        # Model config
+        parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'),
+                            help="Type of positional embedding to use on top of the image features")
+        # * Transformer
+        parser.add_argument('--enc_layers', default=3, type=int,
+                            help="Number of encoding layers in the transformer")
+        parser.add_argument('--dec_layers', default=3, type=int,
+                            help="Number of decoding layers in the transformer")
+        parser.add_argument('--t2v_layers', default=2, type=int,
+                            help="Number of decoding layers in the transformer")
+        parser.add_argument('--sent_layers', default=1, type=int,
+                            help="Number of decoding layers in the transformer")
+        parser.add_argument('--moment_layers', default=1, type=int,
+                            help="Number of decoding layers in the transformer")
+        parser.add_argument('--dummy_layers', default=2, type=int,
+                            help="Number of encoding layers in the transformer")
+        parser.add_argument('--dim_feedforward', default=1024, type=int,
+                            help="Intermediate size of the feedforward layers in the transformer blocks")
+        parser.add_argument('--hidden_dim', default=256, type=int,
+                            help="Size of the embeddings (dimension of the transformer)")
+        parser.add_argument('--input_dropout', default=0.5, type=float,
+                            help="Dropout applied in input")
+        parser.add_argument('--dropout', default=0.1, type=float,
+                            help="Dropout applied in the transformer")
+        parser.add_argument("--txt_drop_ratio", default=0, type=float,
+                            help="drop txt_drop_ratio tokens from text input. 0.1=10%")
+        parser.add_argument("--use_txt_pos", action="store_true", help="use position_embedding for text as well.")
+        parser.add_argument('--nheads', default=8, type=int,
+                            help="Number of attention heads inside the transformer's attentions")
+        parser.add_argument('--num_queries', default=10, type=int,
+                            help="Number of query slots")
+        parser.add_argument('--num_dummies', default=45, type=int,
+                            help="Number of dummy tokens")
+        parser.add_argument('--total_prompts', default=10, type=int,
+                            help="Number of query slots")
+        parser.add_argument('--num_prompts', default=1, type=int,
+                            help="Number of dummy tokens")
+        parser.add_argument('--pre_norm', action='store_true')
+        # other model configs
+        parser.add_argument("--n_input_proj", type=int, default=2, help="#layers to encoder input")
+        parser.add_argument("--contrastive_hdim", type=int, default=64, help="dim for contrastive embeddings")
+        parser.add_argument("--temperature", type=float, default=0.07, help="temperature nce contrastive_align_loss")
+        # Loss
+
+        parser.add_argument("--saliency_margin", type=float, default=0.2)
+        parser.add_argument('--no_aux_loss', dest='aux_loss', action='store_false',
+                            help="Disables auxiliary decoding losses (loss at each layer)")
+        parser.add_argument("--span_loss_type", default="l1", type=str, choices=['l1', 'ce'],
+                            help="l1: (center-x, width) regression. ce: (st_idx, ed_idx) classification.")
+        parser.add_argument("--contrastive_align_loss", action="store_true",
+                            help="Disable contrastive_align_loss between matched query spans and the text.")
+        # * Matcher
+        parser.add_argument('--set_cost_span', default=10, type=float,
+                            help="L1 span coefficient in the matching cost")
+        parser.add_argument('--set_cost_giou', default=1, type=float,
+                            help="giou span coefficient in the matching cost")
+        parser.add_argument('--set_cost_class', default=4, type=float,
+                            help="Class coefficient in the matching cost")
+
+        # * Loss coefficients
+        parser.add_argument("--lw_saliency", type=float, default=1.,
+                            help="weight for saliency loss, set to 0 will ignore")
+        parser.add_argument("--lw_wattn", type=float, default=1.,
+                            help="weight for saliency loss, set to 0 will ignore")
+        parser.add_argument("--lw_ms_align", type=float, default=1.,
+                            help="weight for saliency loss, set to 0 will ignore")
+        parser.add_argument("--lw_distill", type=float, default=1.,
+                            help="weight for saliency loss, set to 0 will ignore")
+        parser.add_argument('--span_loss_coef', default=10, type=float)
+        parser.add_argument('--giou_loss_coef', default=1, type=float)
+        parser.add_argument('--label_loss_coef', default=4, type=float)
+        parser.add_argument('--eos_coef', default=0.1, type=float,
+                            help="Relative classification weight of the no-object class")
+        parser.add_argument("--contrastive_align_loss_coef", default=0.0, type=float)
+
+        parser.add_argument("--no_sort_results", action="store_true",
+                            help="do not sort results, use this for moment query visualization")
+        parser.add_argument("--max_before_nms", type=int, default=10)
+        parser.add_argument("--max_after_nms", type=int, default=10)
+        parser.add_argument("--conf_thd", type=float, default=0.0, help="only keep windows with conf >= conf_thd")
+        parser.add_argument("--nms_thd", type=float, default=-1,
+                            help="additionally use non-maximum suppression "
+                                 "(or non-minimum suppression for distance)"
+                                 "to post-processing the predictions. "
+                                 "-1: do not use nms. [0, 1]")
+        self.parser = parser
+
+    def display_save(self, opt):
+        args = vars(opt)
+        # Display settings
+        print(dict_to_markdown(vars(opt), max_str_len=120))
+        # Save settings
+        if not isinstance(self, TestOptions):
+            option_file_path = os.path.join(opt.results_dir, self.saved_option_filename)  # not yaml file indeed
+            save_json(args, option_file_path, save_pretty=True)
+
+    def parse(self, a_feat_dir=None):
+        if not self.initialized:
+            self.initialize()
+        opt = self.parser.parse_args()
+
+        if opt.debug:
+            opt.results_root = os.path.sep.join(opt.results_root.split(os.path.sep)[:-1] + ["debug_results", ])
+            opt.num_workers = 0
+
+        if isinstance(self, TestOptions):
+            # modify model_dir to absolute path
+            # opt.model_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results", opt.model_dir)
+            opt.model_dir = os.path.dirname(opt.resume)
+            if a_feat_dir is not None:
+                opt.a_feat_dir = a_feat_dir
+            saved_options = load_json(os.path.join(opt.model_dir, self.saved_option_filename))
+            for arg in saved_options:  # use saved options to overwrite all BaseOptions args.
+                if arg not in ["results_root", "num_workers", "nms_thd", "debug",  # "max_before_nms", "max_after_nms"
+                               "max_pred_l", "min_pred_l",
+                               "resume", "resume_all", "no_sort_results"]:
+                    setattr(opt, arg, saved_options[arg])
+            # opt.no_core_driver = True
+            if opt.eval_results_dir is not None:
+                opt.results_dir = opt.eval_results_dir
+        else:
+            if opt.exp_id is None:
+                raise ValueError("--exp_id is required for at a training option!")
+
+            ctx_str = opt.ctx_mode + "_sub" if any(["sub_ctx" in p for p in opt.v_feat_dirs]) else opt.ctx_mode
+            opt.results_dir = os.path.join(opt.results_root,
+                                           "-".join([opt.dset_name, ctx_str, opt.exp_id,
+                                                     str(opt.enc_layers) + str(opt.dec_layers) + str(opt.t2v_layers) + str(opt.moment_layers) + str(opt.dummy_layers) + str(opt.sent_layers),
+                                                     'ndum_' + str(opt.num_dummies), 'nprom_' + str(opt.num_prompts) + '_' + str(opt.total_prompts)]))
+            mkdirp(opt.results_dir)
+            save_fns = ['cg_detr/model.py', 'cg_detr/transformer.py']
+            for save_fn in save_fns:
+                shutil.copyfile(save_fn, os.path.join(opt.results_dir, os.path.basename(save_fn)))
+
+            # save a copy of current code
+            code_dir = os.path.dirname(os.path.realpath(__file__))
+            code_zip_filename = os.path.join(opt.results_dir, "code.zip")
+            make_zipfile(code_dir, code_zip_filename,
+                         enclosing_dir="code",
+                         exclude_dirs_substring="results",
+                         exclude_dirs=["results", "debug_results", "__pycache__"],
+                         exclude_extensions=[".pyc", ".ipynb", ".swap"], )
+
+        self.display_save(opt)
+
+        opt.ckpt_filepath = os.path.join(opt.results_dir, self.ckpt_filename)
+        opt.train_log_filepath = os.path.join(opt.results_dir, self.train_log_filename)
+        opt.eval_log_filepath = os.path.join(opt.results_dir, self.eval_log_filename)
+        opt.tensorboard_log_dir = os.path.join(opt.results_dir, self.tensorboard_log_dir)
+        opt.device = torch.device("cuda" if opt.device >= 0 else "cpu")
+        opt.pin_memory = not opt.no_pin_memory
+
+        opt.use_tef = "tef" in opt.ctx_mode
+        opt.use_video = "video" in opt.ctx_mode
+        if not opt.use_video:
+            opt.v_feat_dim = 0
+        if opt.use_tef:
+            opt.v_feat_dim += 2
+
+        self.opt = opt
+        return opt
+
+
+class TestOptions(BaseOptions):
+    """add additional options for evaluating"""
+
+    def initialize(self):
+        BaseOptions.initialize(self)
+        # also need to specify --eval_split_name
+        self.parser.add_argument("--eval_id", type=str, help="evaluation id")
+        self.parser.add_argument("--eval_results_dir", type=str, default=None,
+                                 help="dir to save results, if not set, fall back to training results_dir")
+        self.parser.add_argument("--model_dir", type=str,
+                                 help="dir contains the model file, will be converted to absolute path afterwards")
+

third_party/cgdetr/cg_detr/crossattention.py

@@ -0,0 +1,396 @@
+# ------------------------------------------------------------------------
+# DAB-DETR
+# Copyright (c) 2022 IDEA. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from Conditional DETR (https://github.com/Atten4Vis/ConditionalDETR)
+# Copyright (c) 2021 Microsoft. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+# ------------------------------------------------------------------------
+# Modified from codes in torch.nn
+# ------------------------------------------------------------------------
+
+"""
+MultiheadAttention that support query, key, and value to have different dimensions.
+Query, key, and value projections are removed.
+Mostly copy-paste from https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/activation.py#L873
+and https://github.com/pytorch/pytorch/blob/master/torch/nn/functional.py#L4837
+"""
+
+import copy
+from typing import Optional, List
+
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+import warnings
+from typing import Tuple, Optional
+
+import torch
+from torch import Tensor
+from torch.nn.modules.linear import Linear
+from torch.nn.init import xavier_uniform_
+from torch.nn.init import constant_
+from torch.nn.init import xavier_normal_
+from torch.nn.parameter import Parameter
+from torch.nn.modules.module import Module
+from torch.nn import functional as F
+
+import warnings
+import math
+
+from torch._C import _infer_size, _add_docstr
+from torch.nn import _reduction as _Reduction
+from torch.nn.modules import utils
+from torch.nn.modules.utils import _single, _pair, _triple, _list_with_default
+from torch.nn import grad
+from torch import _VF
+from torch._jit_internal import boolean_dispatch, List, Optional, _overload, Tuple
+try:
+    from torch.overrides import has_torch_function, handle_torch_function
+except:
+    from torch._overrides import has_torch_function, handle_torch_function
+Tensor = torch.Tensor
+
+from torch.nn.functional import linear, pad, softmax, dropout
+
+class MultiheadAttention(Module):
+    r"""Allows the model to jointly attend to information
+    from different representation subspaces.
+    See reference: Attention Is All You Need
+    .. math::
+        \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O
+        \text{where} head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)
+    Args:
+        embed_dim: total dimension of the model.
+        num_heads: parallel attention heads.
+        dropout: a Dropout layer on attn_output_weights. Default: 0.0.
+        bias: add bias as module parameter. Default: True.
+        add_bias_kv: add bias to the key and value sequences at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        kdim: total number of features in key. Default: None.
+        vdim: total number of features in value. Default: None.
+        Note: if kdim and vdim are None, they will be set to embed_dim such that
+        query, key, and value have the same number of features.
+    Examples::
+        >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads)
+        >>> attn_output, attn_output_weights = multihead_attn(query, key, value)
+    """
+    bias_k: Optional[torch.Tensor]
+    bias_v: Optional[torch.Tensor]
+
+    def __init__(self, embed_dim, num_heads, dropout=0., num_dummies=3, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None):
+        super(MultiheadAttention, self).__init__()
+        self.num_dummies = num_dummies
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        vdim = vdim if vdim is not None else embed_dim
+        self.out_proj = Linear(vdim , vdim)
+
+        self.in_proj_bias = None
+        self.in_proj_weight = None
+        self.bias_k = self.bias_v = None
+        self.q_proj_weight = None
+        self.k_proj_weight = None
+        self.v_proj_weight = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self._reset_parameters()
+
+    def _reset_parameters(self):
+        constant_(self.out_proj.bias, 0.)
+
+    def __setstate__(self, state):
+        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
+        if '_qkv_same_embed_dim' not in state:
+            state['_qkv_same_embed_dim'] = True
+
+        super(MultiheadAttention, self).__setstate__(state)
+
+    def forward(self, query, key, value, key_padding_mask=None,
+                need_weights=True, attn_mask=None, dummy=True):
+        # type: (Tensor, Tensor, Tensor, Optional[Tensor], bool, Optional[Tensor]) -> Tuple[Tensor, Optional[Tensor]]
+        r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. When given a binary mask and a value is True,
+            the corresponding value on the attention layer will be ignored. When given
+            a byte mask and a value is non-zero, the corresponding value on the attention
+            layer will be ignored
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+    Shape:
+        - Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the position
+          with the zero positions will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*\text{num_heads}, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensure that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          is not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+        """
+        if not self._qkv_same_embed_dim:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
+                v_proj_weight=self.v_proj_weight, out_dim=self.vdim, num_dummies=self.num_dummies, dummy=dummy)
+        else:
+            return multi_head_attention_forward(
+                query, key, value, self.embed_dim, self.num_heads,
+                self.in_proj_weight, self.in_proj_bias,
+                self.bias_k, self.bias_v, self.add_zero_attn,
+                self.dropout, self.out_proj.weight, self.out_proj.bias,
+                training=self.training,
+                key_padding_mask=key_padding_mask, need_weights=need_weights,
+                attn_mask=attn_mask, out_dim=self.vdim, num_dummies=self.num_dummies, dummy=dummy)
+
+
+def multi_head_attention_forward(query: Tensor,
+                                 key: Tensor,
+                                 value: Tensor,
+                                 embed_dim_to_check: int,
+                                 num_heads: int,
+                                 in_proj_weight: Tensor,
+                                 in_proj_bias: Tensor,
+                                 bias_k: Optional[Tensor],
+                                 bias_v: Optional[Tensor],
+                                 add_zero_attn: bool,
+                                 dropout_p: float,
+                                 out_proj_weight: Tensor,
+                                 out_proj_bias: Tensor,
+                                 training: bool = True,
+                                 key_padding_mask: Optional[Tensor] = None,
+                                 need_weights: bool = True,
+                                 attn_mask: Optional[Tensor] = None,
+                                 use_separate_proj_weight: bool = False,
+                                 q_proj_weight: Optional[Tensor] = None,
+                                 k_proj_weight: Optional[Tensor] = None,
+                                 v_proj_weight: Optional[Tensor] = None,
+                                 static_k: Optional[Tensor] = None,
+                                 static_v: Optional[Tensor] = None,
+                                 out_dim: Optional[Tensor] = None,
+                                 num_dummies=3,
+                                 dummy=True,
+                                 ) -> Tuple[Tensor, Optional[Tensor]]:
+    r"""
+    Args:
+        query, key, value: map a query and a set of key-value pairs to an output.
+            See "Attention Is All You Need" for more details.
+        embed_dim_to_check: total dimension of the model.
+        num_heads: parallel attention heads.
+        in_proj_weight, in_proj_bias: input projection weight and bias.
+        bias_k, bias_v: bias of the key and value sequences to be added at dim=0.
+        add_zero_attn: add a new batch of zeros to the key and
+                       value sequences at dim=1.
+        dropout_p: probability of an element to be zeroed.
+        out_proj_weight, out_proj_bias: the output projection weight and bias.
+        training: apply dropout if is ``True``.
+        key_padding_mask: if provided, specified padding elements in the key will
+            be ignored by the attention. This is an binary mask. When the value is True,
+            the corresponding value on the attention layer will be filled with -inf.
+        need_weights: output attn_output_weights.
+        attn_mask: 2D or 3D mask that prevents attention to certain positions. A 2D mask will be broadcasted for all
+            the batches while a 3D mask allows to specify a different mask for the entries of each batch.
+        use_separate_proj_weight: the function accept the proj. weights for query, key,
+            and value in different forms. If false, in_proj_weight will be used, which is
+            a combination of q_proj_weight, k_proj_weight, v_proj_weight.
+        q_proj_weight, k_proj_weight, v_proj_weight, in_proj_bias: input projection weight and bias.
+        static_k, static_v: static key and value used for attention operators.
+    Shape:
+        Inputs:
+        - query: :math:`(L, N, E)` where L is the target sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key: :math:`(S, N, E)`, where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - value: :math:`(S, N, E)` where S is the source sequence length, N is the batch size, E is
+          the embedding dimension.
+        - key_padding_mask: :math:`(N, S)` where N is the batch size, S is the source sequence length.
+          If a ByteTensor is provided, the non-zero positions will be ignored while the zero positions
+          will be unchanged. If a BoolTensor is provided, the positions with the
+          value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+          3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+          S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+          positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
+          while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``True``
+          are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+          is provided, it will be added to the attention weight.
+        - static_k: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        - static_v: :math:`(N*num_heads, S, E/num_heads)`, where S is the source sequence length,
+          N is the batch size, E is the embedding dimension. E/num_heads is the head dimension.
+        Outputs:
+        - attn_output: :math:`(L, N, E)` where L is the target sequence length, N is the batch size,
+          E is the embedding dimension.
+        - attn_output_weights: :math:`(N, L, S)` where N is the batch size,
+          L is the target sequence length, S is the source sequence length.
+    """
+    if not torch.jit.is_scripting():
+        tens_ops = (query, key, value, in_proj_weight, in_proj_bias, bias_k, bias_v,
+                    out_proj_weight, out_proj_bias)
+        if any([type(t) is not Tensor for t in tens_ops]) and has_torch_function(tens_ops):
+            return handle_torch_function(
+                multi_head_attention_forward, tens_ops, query, key, value,
+                embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias,
+                bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight,
+                out_proj_bias, training=training, key_padding_mask=key_padding_mask,
+                need_weights=need_weights, attn_mask=attn_mask,
+                use_separate_proj_weight=use_separate_proj_weight,
+                q_proj_weight=q_proj_weight, k_proj_weight=k_proj_weight,
+                v_proj_weight=v_proj_weight, static_k=static_k, static_v=static_v)
+    tgt_len, bsz, embed_dim = query.size()
+    assert embed_dim == embed_dim_to_check
+    # allow MHA to have different sizes for the feature dimension
+    assert key.size(0) == value.size(0) and key.size(1) == value.size(1)
+
+    head_dim = embed_dim // num_heads
+    v_head_dim = out_dim // num_heads
+    assert head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+    scaling = float(head_dim) ** -0.5
+
+    q = query * scaling
+    k = key
+    v = value
+
+    if attn_mask is not None:
+        assert attn_mask.dtype == torch.float32 or attn_mask.dtype == torch.float64 or \
+            attn_mask.dtype == torch.float16 or attn_mask.dtype == torch.uint8 or attn_mask.dtype == torch.bool, \
+            'Only float, byte, and bool types are supported for attn_mask, not {}'.format(attn_mask.dtype)
+        if attn_mask.dtype == torch.uint8:
+            warnings.warn("Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+            attn_mask = attn_mask.to(torch.bool)
+
+        if attn_mask.dim() == 2:
+            attn_mask = attn_mask.unsqueeze(0)
+            if list(attn_mask.size()) != [1, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 2D attn_mask is not correct.')
+        elif attn_mask.dim() == 3:
+            if list(attn_mask.size()) != [bsz * num_heads, query.size(0), key.size(0)]:
+                raise RuntimeError('The size of the 3D attn_mask is not correct.')
+        else:
+            raise RuntimeError("attn_mask's dimension {} is not supported".format(attn_mask.dim()))
+        # attn_mask's dim is 3 now.
+
+    # convert ByteTensor key_padding_mask to bool
+    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:
+        warnings.warn("Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.")
+        key_padding_mask = key_padding_mask.to(torch.bool)
+
+    if bias_k is not None and bias_v is not None:
+        if static_k is None and static_v is None:
+            k = torch.cat([k, bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = pad(attn_mask, (0, 1))
+            if key_padding_mask is not None:
+                key_padding_mask = pad(key_padding_mask, (0, 1))
+        else:
+            assert static_k is None, "bias cannot be added to static key."
+            assert static_v is None, "bias cannot be added to static value."
+    else:
+        assert bias_k is None
+        assert bias_v is None
+
+    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
+    if k is not None:
+        k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)
+    if v is not None:
+        v = v.contiguous().view(-1, bsz * num_heads, v_head_dim).transpose(0, 1)
+
+    if static_k is not None:
+        assert static_k.size(0) == bsz * num_heads
+        assert static_k.size(2) == head_dim
+        k = static_k
+
+    if static_v is not None:
+        assert static_v.size(0) == bsz * num_heads
+        assert static_v.size(2) == v_head_dim
+        v = static_v
+
+    src_len = k.size(1)
+
+    if key_padding_mask is not None:
+        assert key_padding_mask.size(0) == bsz
+        assert key_padding_mask.size(1) == src_len
+
+    if add_zero_attn:
+        src_len += 1
+        k = torch.cat([k, torch.zeros((k.size(0), 1) + k.size()[2:], dtype=k.dtype, device=k.device)], dim=1)
+        v = torch.cat([v, torch.zeros((v.size(0), 1) + v.size()[2:], dtype=v.dtype, device=v.device)], dim=1)
+        if attn_mask is not None:
+            attn_mask = pad(attn_mask, (0, 1))
+        if key_padding_mask is not None:
+            key_padding_mask = pad(key_padding_mask, (0, 1))
+
+    attn_output_weights = torch.bmm(q, k.transpose(1, 2))
+    assert list(attn_output_weights.size()) == [bsz * num_heads, tgt_len, src_len]
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_output_weights.masked_fill_(attn_mask, float('-inf'))
+        else:
+            attn_output_weights += attn_mask
+
+
+    if key_padding_mask is not None:
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        attn_output_weights = attn_output_weights.masked_fill(
+            key_padding_mask.unsqueeze(1).unsqueeze(2),
+            float('-inf'),
+        )
+        attn_output_weights = attn_output_weights.view(bsz * num_heads, tgt_len, src_len)
+
+    attn_output_weights = softmax(attn_output_weights, dim=-1)
+    attn_output_weights_d = dropout(attn_output_weights, p=dropout_p, training=training)
+    if dummy:
+        attn_output = torch.bmm(attn_output_weights_d[:, :, num_dummies:], v[:, num_dummies:,:])
+    else:
+        attn_output = torch.bmm(attn_output_weights_d, v)
+    assert list(attn_output.size()) == [bsz * num_heads, tgt_len, v_head_dim]
+    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, out_dim)
+    attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
+
+    if need_weights:
+        # average attention weights over heads
+        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)
+        return attn_output, attn_output_weights.sum(dim=1) / num_heads
+    else:
+        return attn_output, None

third_party/cgdetr/cg_detr/inference.py

@@ -0,0 +1,480 @@
+import pprint
+from tqdm import tqdm, trange
+import numpy as np
+import os
+from collections import OrderedDict, defaultdict
+from utils.basic_utils import AverageMeter
+
+import torch
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+from torch.utils.data import DataLoader
+
+from cg_detr.config import TestOptions
+from cg_detr.model import build_model
+from cg_detr.span_utils import span_cxw_to_xx
+from cg_detr.start_end_dataset import StartEndDataset, start_end_collate, prepare_batch_inputs
+from cg_detr.postprocessing_cg_detr import PostProcessorDETR
+from standalone_eval.eval import eval_submission
+from utils.basic_utils import save_jsonl, save_json
+from utils.temporal_nms import temporal_nms
+
+import logging
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(format="%(asctime)s.%(msecs)03d:%(levelname)s:%(name)s - %(message)s",
+                    datefmt="%Y-%m-%d %H:%M:%S",
+                    level=logging.INFO)
+
+
+def post_processing_mr_nms(mr_res, nms_thd, max_before_nms, max_after_nms):
+    mr_res_after_nms = []
+    for e in mr_res:
+        e["pred_relevant_windows"] = temporal_nms(
+            e["pred_relevant_windows"][:max_before_nms],
+            nms_thd=nms_thd,
+            max_after_nms=max_after_nms
+        )
+        mr_res_after_nms.append(e)
+    return mr_res_after_nms
+
+
+def eval_epoch_post_processing(submission, opt, gt_data, save_submission_filename):
+    # IOU_THDS = (0.5, 0.7)
+    logger.info("Saving/Evaluating before nms results")
+    submission_path = os.path.join(opt.results_dir, save_submission_filename)
+    save_jsonl(submission, submission_path)
+
+    if opt.eval_split_name in ["val"]:  # since test_public has no GT
+        metrics = eval_submission(
+            submission, gt_data,
+            verbose=opt.debug, match_number=not opt.debug
+        )
+        save_metrics_path = submission_path.replace(".jsonl", "_metrics.json")
+        save_json(metrics, save_metrics_path, save_pretty=True, sort_keys=False)
+        latest_file_paths = [submission_path, save_metrics_path]
+    else:
+        metrics = None
+        latest_file_paths = [submission_path, ]
+
+    if opt.nms_thd != -1:
+        logger.info("[MR] Performing nms with nms_thd {}".format(opt.nms_thd))
+        submission_after_nms = post_processing_mr_nms(
+            submission, nms_thd=opt.nms_thd,
+            max_before_nms=opt.max_before_nms, max_after_nms=opt.max_after_nms
+        )
+
+        logger.info("Saving/Evaluating nms results")
+        submission_nms_path = submission_path.replace(".jsonl", "_nms_thd_{}.jsonl".format(opt.nms_thd))
+        save_jsonl(submission_after_nms, submission_nms_path)
+        if opt.eval_split_name == "val":
+            metrics_nms = eval_submission(
+                submission_after_nms, gt_data,
+                verbose=opt.debug, match_number=not opt.debug
+            )
+            save_metrics_nms_path = submission_nms_path.replace(".jsonl", "_metrics.json")
+            save_json(metrics_nms, save_metrics_nms_path, save_pretty=True, sort_keys=False)
+            latest_file_paths += [submission_nms_path, save_metrics_nms_path]
+        else:
+            metrics_nms = None
+            latest_file_paths = [submission_nms_path, ]
+    else:
+        metrics_nms = None
+    return metrics, metrics_nms, latest_file_paths
+
+
+# for HL
+@torch.no_grad()
+def compute_hl_results(model, eval_loader, opt, epoch_i=None, criterion=None, tb_writer=None):
+    model.eval()
+    if criterion:
+        assert eval_loader.dataset.load_labels
+        criterion.eval()
+
+    loss_meters = defaultdict(AverageMeter)
+    write_tb = tb_writer is not None and epoch_i is not None
+
+    mr_res = []
+
+    topk = 5 # top-5 map
+
+    video_ap_collected = []
+    for batch in tqdm(eval_loader, desc="compute st ed scores"):
+        query_meta = batch[0]
+
+        model_inputs, targets = prepare_batch_inputs(batch[1], opt.device, non_blocking=opt.pin_memory)
+
+        outputs = model(**model_inputs)
+
+        # loss meters
+        # if criterion:
+        #     loss_dict = criterion(outputs, targets)
+        #     weight_dict = criterion.weight_dict
+            # print(loss_dict)
+            # print(weight_dict)
+            # print('#######')
+            # {'loss_saliency': tensor(18.1374, device='cuda:0')}
+            # {'loss_span': 10, 'loss_giou': 1, 'loss_label': 4, 'loss_saliency': 1.0, 'loss_ms_align': 1.0,
+            #  'loss_distill': 1.0, 'loss_span_0': 10, 'loss_giou_0': 1, 'loss_label_0': 4, 'loss_ms_align_0': 1.0,
+            #  'loss_distill_0': 1.0}
+            # losses=0.
+            # print(loss_dict.keys(), weight_dict.keys())
+            # losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
+            # loss_dict["loss_overall"] = float(losses)  # for logging only
+            # print(loss_dict.items())
+            #
+            # print(weight_dict.items())
+            # for k, v in loss_dict.items():
+            #     loss_meters[k].update(float(v) * weight_dict[k] if k in weight_dict else float(v))
+
+
+        preds = outputs['saliency_scores'].clone().detach()
+
+        for meta, pred in zip(query_meta, preds):
+            pred = pred
+            label = meta['label'] # raw label
+
+            video_ap = []
+            # Follow the UMT code "https://github.com/TencentARC/UMT/blob/main/datasets/tvsum.py"
+            
+            if opt.dset_name in ["tvsum"]:
+                for i in range(20):
+                    pred=pred.cpu()
+                    cur_pred = pred[:len(label)]
+                    inds = torch.argsort(cur_pred, descending=True, dim=-1)
+
+                    # video_id = self.get_video_id(idx)
+                    cur_label = torch.Tensor(label)[:, i]
+                    cur_label = torch.where(cur_label > cur_label.median(), 1.0, .0)
+
+                    cur_label = cur_label[inds].tolist()[:topk]
+
+                    # if (num_gt := sum(cur_label)) == 0:
+                    num_gt = sum(cur_label)
+                    if num_gt == 0:
+                        video_ap.append(0)
+                        continue
+
+                    hits = ap = rec = 0
+                    prc = 1
+
+                    for j, gt in enumerate(cur_label):
+                        hits += gt
+
+                        _rec = hits / num_gt
+                        _prc = hits / (j + 1)
+
+                        ap += (_rec - rec) * (prc + _prc) / 2
+                        rec, prc = _rec, _prc
+
+                    video_ap.append(ap)
+            
+            elif opt.dset_name in ["youtube_uni"]:
+                cur_pred = pred[:len(label)]
+                # if opt.dset_name == "tvsum_sfc":
+                cur_pred = cur_pred.cpu()
+                inds = torch.argsort(cur_pred, descending=True, dim=-1)
+
+
+                cur_label = torch.Tensor(label).squeeze()[inds].tolist()
+                
+                num_gt = sum(cur_label)
+                if num_gt == 0:
+                    video_ap.append(0)
+                    continue
+
+                hits = ap = rec = 0
+                prc = 1
+
+                for j, gt in enumerate(cur_label):
+                    hits += gt
+
+                    _rec = hits / num_gt
+                    _prc = hits / (j + 1)
+
+                    ap += (_rec - rec) * (prc + _prc) / 2
+                    rec, prc = _rec, _prc
+                
+                video_ap.append(float(ap))
+            else:
+                print("No such dataset")
+                exit(-1)
+                    
+            video_ap_collected.append(video_ap)
+
+    mean_ap = np.mean(video_ap_collected)
+    submmission = dict(mAP=round(mean_ap, 5))
+    
+
+    # tensorboard writer
+    if write_tb and criterion:
+        for k, v in loss_meters.items():
+            tb_writer.add_scalar("Eval/{}".format(k), v.avg, epoch_i + 1)
+
+    return submmission, loss_meters 
+
+
+
+@torch.no_grad()
+def compute_mr_results(model, eval_loader, opt, epoch_i=None, criterion=None, tb_writer=None):
+    model.eval()
+    if criterion:
+        assert eval_loader.dataset.load_labels
+        criterion.eval()
+
+    loss_meters = defaultdict(AverageMeter)
+    write_tb = tb_writer is not None and epoch_i is not None
+
+    mr_res = []
+    for batch in tqdm(eval_loader, desc="compute st ed scores"):
+        query_meta = batch[0]
+
+        model_inputs, targets = prepare_batch_inputs(batch[1], opt.device, non_blocking=opt.pin_memory)
+
+        outputs = model(**model_inputs)
+        prob = F.softmax(outputs["pred_logits"], -1)  # (batch_size, #queries, #classes=2)
+        if opt.span_loss_type == "l1":
+            scores = prob[..., 0]  # * (batch_size, #queries)  foreground label is 0, we directly take it
+            pred_spans = outputs["pred_spans"]  # (bsz, #queries, 2)
+            _saliency_scores = outputs["saliency_scores"].half()  # (bsz, L)
+            saliency_scores = []
+            valid_vid_lengths = model_inputs["src_vid_mask"].sum(1).cpu().tolist()
+            for j in range(len(valid_vid_lengths)):
+                saliency_scores.append(_saliency_scores[j, :int(valid_vid_lengths[j])].tolist())
+        else:
+            bsz, n_queries = outputs["pred_spans"].shape[:2]  # # (bsz, #queries, max_v_l *2)
+            pred_spans_logits = outputs["pred_spans"].view(bsz, n_queries, 2, opt.max_v_l)
+            pred_span_scores, pred_spans = F.softmax(pred_spans_logits, dim=-1).max(-1)  # 2 * (bsz, #queries, 2)
+            scores = torch.prod(pred_span_scores, 2)  # (bsz, #queries)
+            pred_spans[:, 1] += 1
+            pred_spans *= opt.clip_length
+
+        # compose predictions
+        for idx, (meta, spans, score) in enumerate(zip(query_meta, pred_spans.cpu(), scores.cpu())):
+            if opt.span_loss_type == "l1":
+                spans = span_cxw_to_xx(spans) * meta["duration"]
+                spans = torch.clamp(spans, 0, meta["duration"])
+            # # (#queries, 3), [st(float), ed(float), score(float)]
+            cur_ranked_preds = torch.cat([spans, score[:, None]], dim=1).tolist()
+            if not opt.no_sort_results:
+                cur_ranked_preds = sorted(cur_ranked_preds, key=lambda x: x[2], reverse=True)
+            cur_ranked_preds = [[float(f"{e:.4f}") for e in row] for row in cur_ranked_preds]
+            cur_query_pred = dict(
+                qid=meta["qid"],
+                query=meta["query"],
+                vid=meta["vid"],
+                pred_relevant_windows=cur_ranked_preds,
+                pred_saliency_scores=saliency_scores[idx]
+            )
+            mr_res.append(cur_query_pred)
+
+        if criterion:
+            loss_dict = criterion(outputs, targets)
+            weight_dict = criterion.weight_dict
+            losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
+            loss_dict["loss_overall"] = float(losses)  # for logging only
+            for k, v in loss_dict.items():
+                loss_meters[k].update(float(v) * weight_dict[k] if k in weight_dict else float(v))
+
+        if opt.debug:
+            break
+
+    if write_tb and criterion:
+        for k, v in loss_meters.items():
+            tb_writer.add_scalar("Eval/{}".format(k), v.avg, epoch_i + 1)
+
+    if opt.dset_name in ['hl']:
+        post_processor = PostProcessorDETR(
+            clip_length=opt.clip_length, min_ts_val=0, max_ts_val=150,
+            min_w_l=2, max_w_l=150, move_window_method="left",
+            process_func_names=("clip_ts", "round_multiple")
+        )
+    elif opt.dset_name in ['charadesSTA']:
+        if opt.v_feat_dim == 4096: # vgg
+            post_processor = PostProcessorDETR(
+                clip_length=opt.clip_length, min_ts_val=0, max_ts_val=360,
+                min_w_l=12, max_w_l=360, move_window_method="left",
+                process_func_names=("clip_ts", "round_multiple")
+            )
+        else:
+            post_processor = PostProcessorDETR(
+                clip_length=opt.clip_length, min_ts_val=0, max_ts_val=150,
+                min_w_l=2, max_w_l=60, move_window_method="left",
+                process_func_names=("clip_ts", "round_multiple")
+            )
+    else:
+        post_processor = PostProcessorDETR(
+            clip_length=opt.clip_length, min_ts_val=0, max_ts_val=50000,
+            min_w_l=0, max_w_l=50000, move_window_method="left",
+            process_func_names=(["round_multiple"])
+        )
+
+    mr_res = post_processor(mr_res)
+    return mr_res, loss_meters
+
+
+def get_eval_res(model, eval_loader, opt, epoch_i, criterion, tb_writer):
+    """compute and save query and video proposal embeddings"""
+    eval_res, eval_loss_meters = compute_mr_results(model, eval_loader, opt, epoch_i, criterion, tb_writer)  # list(dict)
+    return eval_res, eval_loss_meters
+
+
+def eval_epoch(model, eval_dataset, opt, save_submission_filename, epoch_i=None, criterion=None, tb_writer=None):
+    logger.info("Generate submissions")
+    model.eval()
+    if criterion is not None and eval_dataset.load_labels:
+        criterion.eval()
+    else:
+        criterion = None
+
+    if opt.dset_name == 'tacos':
+        shuffle = True
+    else:
+        shuffle = False
+
+    eval_loader = DataLoader(
+        eval_dataset,
+        collate_fn=start_end_collate,
+        batch_size=opt.eval_bsz,
+        num_workers=opt.num_workers,
+        shuffle=shuffle,
+        pin_memory=opt.pin_memory
+    )
+
+
+    # tvsum 
+    if opt.dset_name in ['tvsum', 'youtube_uni']:
+        metrics, eval_loss_meters = compute_hl_results(model, eval_loader, opt, epoch_i, criterion, tb_writer)
+        
+        # to match original save format
+        submission = [
+            {"brief": metrics}
+        ]
+        submission_path = os.path.join(opt.results_dir, "latest_metric.jsonl")
+        save_jsonl(submission, submission_path)
+
+        return submission[0], submission[0], eval_loss_meters, [submission_path]
+
+    else:
+        submission, eval_loss_meters = get_eval_res(model, eval_loader, opt, epoch_i, criterion, tb_writer)
+
+        if opt.dset_name in ['charadesSTA', 'tacos', 'nlq']:
+            new_submission = []
+            for s in submission:
+                s.pop('pred_saliency_scores', None)
+                new_submission.append(s)
+            submission = new_submission
+
+        if opt.no_sort_results:
+            save_submission_filename = save_submission_filename.replace(".jsonl", "_unsorted.jsonl")
+        metrics, metrics_nms, latest_file_paths = eval_epoch_post_processing(
+            submission, opt, eval_dataset.data, save_submission_filename)
+        return metrics, metrics_nms, eval_loss_meters, latest_file_paths
+
+
+def setup_model(opt):
+    """setup model/optimizer/scheduler and load checkpoints when needed"""
+    logger.info("setup model/optimizer/scheduler")
+    model, criterion = build_model(opt)
+    if opt.device.type == "cuda":
+        logger.info("CUDA enabled.")
+        model.to(opt.device)
+        criterion.to(opt.device)
+
+    param_dicts = [{"params": [p for n, p in model.named_parameters() if p.requires_grad]}]
+    optimizer = torch.optim.AdamW(param_dicts, lr=opt.lr, weight_decay=opt.wd)
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, opt.lr_drop)
+
+    if opt.resume is not None:
+        logger.info(f"Load checkpoint from {opt.resume}")
+        checkpoint = torch.load(opt.resume, map_location="cpu")
+        from collections import OrderedDict
+        new_state_dict = OrderedDict()
+        if 'pt' in opt.resume[:-4]:
+            if 'asr' in opt.resume[:25]:
+                model.load_state_dict(checkpoint["model"])
+            else:
+                for k, v in checkpoint["model"].items():
+                    name = k[7:]  # remove `module.`
+                    new_state_dict[name] = v
+                # model.load_state_dict(checkpoint["model"])
+                model.load_state_dict(new_state_dict)
+        else:
+            model.load_state_dict(checkpoint["model"])
+        if opt.resume_all:
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
+            opt.start_epoch = checkpoint['epoch'] + 1
+        logger.info(f"Loaded model saved at epoch {checkpoint['epoch']} from checkpoint: {opt.resume}")
+    else:
+        logger.warning("If you intend to evaluate the model, please specify --resume with ckpt path")
+
+    return model, criterion, optimizer, lr_scheduler
+
+
+def start_inference(train_opt=None, split=None, splitfile=None):
+    if train_opt is not None:
+        opt = TestOptions().parse(train_opt.a_feat_dir)
+    else:
+        opt = TestOptions().parse()
+    if split is not None:
+        opt.eval_split_name = split
+    if splitfile is not None:
+        opt.eval_path = splitfile
+
+    print(opt.eval_split_name)
+    print(opt.eval_path)
+    logger.info("Setup config, data and model...")
+
+
+    cudnn.benchmark = True
+    cudnn.deterministic = False
+
+    assert opt.eval_path is not None
+    if opt.eval_split_name == 'val':
+        loadlabel = True
+    else:
+        loadlabel = False
+
+    eval_dataset = StartEndDataset(
+        dset_name=opt.dset_name,
+        data_path=opt.eval_path,
+        v_feat_dirs=opt.v_feat_dirs,
+        q_feat_dir=opt.t_feat_dir,
+        q_feat_type="last_hidden_state",
+        max_q_l=opt.max_q_l,
+        max_v_l=opt.max_v_l,
+        ctx_mode=opt.ctx_mode,
+        data_ratio=opt.data_ratio,
+        normalize_v=not opt.no_norm_vfeat,
+        normalize_t=not opt.no_norm_tfeat,
+        clip_len=opt.clip_length,
+        max_windows=opt.max_windows,
+        load_labels=loadlabel,  # opt.eval_split_name == "val",
+        span_loss_type=opt.span_loss_type,
+        txt_drop_ratio=0,
+        dset_domain=opt.dset_domain,
+    )
+
+
+
+    model, criterion, _, _ = setup_model(opt)
+
+    save_submission_filename = "hl_{}_submission.jsonl".format(
+        opt.eval_split_name)
+    # save_submission_filename = "inference_{}_{}_{}_preds.jsonl".format(
+    #     opt.dset_name, opt.eval_split_name, opt.eval_id)
+    logger.info("Starting inference...")
+    with torch.no_grad():
+        metrics_no_nms, metrics_nms, eval_loss_meters, latest_file_paths = \
+            eval_epoch(model, eval_dataset, opt, save_submission_filename, criterion=criterion)
+    if opt.eval_split_name == 'val':
+        logger.info("metrics_no_nms {}".format(pprint.pformat(metrics_no_nms["brief"], indent=4)))
+    if metrics_nms is not None:
+        logger.info("metrics_nms {}".format(pprint.pformat(metrics_nms["brief"], indent=4)))
+
+from sys import argv
+if __name__ == '__main__':
+    _,_,_,_,split,_,splitfile = argv
+
+    start_inference(split=split, splitfile=splitfile)

third_party/cgdetr/cg_detr/matcher.py

@@ -0,0 +1,109 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Modules to compute the matching cost and solve the corresponding LSAP.
+"""
+import torch
+from scipy.optimize import linear_sum_assignment
+from torch import nn
+import torch.nn.functional as F
+from third_party.cgdetr.cg_detr.span_utils import generalized_temporal_iou, span_cxw_to_xx
+
+
+class HungarianMatcher(nn.Module):
+    """This class computes an assignment between the targets and the predictions of the network
+
+    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
+    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
+    while the others are un-matched (and thus treated as non-objects).
+    """
+    def __init__(self,  cost_class: float = 1, cost_span: float = 1, cost_giou: float = 1,
+                 span_loss_type: str = "l1", max_v_l: int = 75):
+        """Creates the matcher
+
+        Params:
+            cost_span: This is the relative weight of the L1 error of the span coordinates in the matching cost
+            cost_giou: This is the relative weight of the giou loss of the spans in the matching cost
+        """
+        super().__init__()
+        self.cost_class = cost_class
+        self.cost_span = cost_span
+        self.cost_giou = cost_giou
+        self.span_loss_type = span_loss_type
+        self.max_v_l = max_v_l
+        self.foreground_label = 0
+        assert cost_class != 0 or cost_span != 0 or cost_giou != 0, "all costs cant be 0"
+
+    @torch.no_grad()
+    def forward(self, outputs, targets):
+        """ Performs the matching
+
+        Params:
+            outputs: This is a dict that contains at least these entries:
+                 "pred_spans": Tensor of dim [batch_size, num_queries, 2] with the predicted span coordinates,
+                    in normalized (cx, w) format
+                 ""pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
+
+            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
+                 "spans": Tensor of dim [num_target_spans, 2] containing the target span coordinates. The spans are
+                    in normalized (cx, w) format
+
+        Returns:
+            A list of size batch_size, containing tuples of (index_i, index_j) where:
+                - index_i is the indices of the selected predictions (in order)
+                - index_j is the indices of the corresponding selected targets (in order)
+            For each batch element, it holds:
+                len(index_i) = len(index_j) = min(num_queries, num_target_spans)
+        """
+        bs, num_queries = outputs["pred_spans"].shape[:2]
+        targets = targets["span_labels"]
+        # import pdb; pdb.set_trace()
+
+        # Also concat the target labels and spans
+        out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
+        tgt_spans = torch.cat([v["spans"] for v in targets])  # [num_target_spans in batch, 2]
+        tgt_ids = torch.full([len(tgt_spans)], self.foreground_label)   # [total #spans in the batch]
+
+        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
+        # but approximate it in 1 - prob[target class].
+        # The 1 is a constant that doesn't change the matching, it can be omitted.
+        cost_class = -out_prob[:, tgt_ids]  # [batch_size * num_queries, total #spans in the batch]
+
+        if self.span_loss_type == "l1":
+            # We flatten to compute the cost matrices in a batch
+            out_spans = outputs["pred_spans"].flatten(0, 1)  # [batch_size * num_queries, 2]
+
+            # Compute the L1 cost between spans
+            cost_span = torch.cdist(out_spans.type(torch.float32), tgt_spans.type(torch.float32), p=1)  # [batch_size * num_queries, total #spans in the batch]
+            cost_span = cost_span.type(torch.bfloat16)
+
+            # Compute the giou cost between spans
+            # [batch_size * num_queries, total #spans in the batch]
+            cost_giou = - generalized_temporal_iou(span_cxw_to_xx(out_spans), span_cxw_to_xx(tgt_spans))
+        else:
+            pred_spans = outputs["pred_spans"]  # (bsz, #queries, max_v_l * 2)
+            pred_spans = pred_spans.view(bs * num_queries, 2, self.max_v_l).softmax(-1)  # (bsz * #queries, 2, max_v_l)
+            cost_span = - pred_spans[:, 0][:, tgt_spans[:, 0]] - \
+                pred_spans[:, 1][:, tgt_spans[:, 1]]  # (bsz * #queries, #spans)
+            # pred_spans = pred_spans.repeat(1, n_spans, 1, 1).flatten(0, 1)  # (bsz * #queries * #spans, max_v_l, 2)
+            # tgt_spans = tgt_spans.view(1, n_spans, 2).repeat(bs * num_queries, 1, 1).flatten(0, 1)  # (bsz * #queries * #spans, 2)
+            # cost_span = pred_spans[tgt_spans]
+            # cost_span = cost_span.view(bs * num_queries, n_spans)
+
+            # giou
+            cost_giou = 0
+
+        # Final cost matrix
+        # import ipdb; ipdb.set_trace()
+        C = self.cost_span * cost_span + self.cost_giou * cost_giou + self.cost_class * cost_class
+        C = C.view(bs, num_queries, -1).cpu()
+
+        sizes = [len(v["spans"]) for v in targets]
+        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
+        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
+
+
+def build_matcher(args):
+    return HungarianMatcher(
+        cost_span=args.set_cost_span, cost_giou=args.set_cost_giou,
+        cost_class=args.set_cost_class, span_loss_type=args.span_loss_type, max_v_l=args.max_v_l
+    )

third_party/cgdetr/cg_detr/misc.py

@@ -0,0 +1,21 @@
+import torch
+
+
+@torch.no_grad()
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k
+    output: (#items, #classes)
+    target: int,
+    """
+    maxk = max(topk)
+    num_items = output.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target)
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / num_items))
+    return res

third_party/cgdetr/cg_detr/model.py

@@ -0,0 +1,1178 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+CG-DETR model and criterion classes.
+"""
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from third_party.cgdetr.cg_detr.span_utils import generalized_temporal_iou, span_cxw_to_xx
+
+from third_party.cgdetr.cg_detr.matcher import build_matcher
+from third_party.cgdetr.cg_detr.transformer import build_transformer, TransformerEncoderLayer, TransformerEncoder
+from third_party.cgdetr.cg_detr.position_encoding import build_position_encoding
+from third_party.cgdetr.cg_detr.misc import accuracy
+import numpy as np
+import copy
+
+def inverse_sigmoid(x, eps=1e-3):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1/x2)
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+    if isinstance(module, nn.Linear) and module.bias is not None:
+        module.bias.data.zero_()
+
+def find_nth(vid, underline, n):
+    max_len = len(vid)
+    start = vid.find(underline)
+    while start >= 0 and n > 1:
+        start = vid.find(underline, start+len(underline))
+        n -= 1
+    if start == -1:
+        start = max_len
+    return start
+
+def element_wise_list_equal(listA, listB):
+    res = []
+    for a, b in zip(listA, listB):
+        if a==b:
+            res.append(True)
+        else:
+            res.append(False)
+    return res
+
+class CGDETR(nn.Module):
+    """ CG DETR. """
+
+    def __init__(self, transformer, position_embed, txt_position_embed, txt_dim, vid_dim,
+                 num_queries, input_dropout, aux_loss=False,
+                 contrastive_align_loss=False, contrastive_hdim=64,
+                 max_v_l=75, span_loss_type="l1", use_txt_pos=False, n_input_proj=2, aud_dim=0, args=None):
+        """ Initializes the model.
+        Parameters:
+            transformer: torch module of the transformer architecture. See transformer.py
+            position_embed: torch module of the position_embedding, See position_encoding.py
+            txt_position_embed: position_embedding for text
+            txt_dim: int, text query input dimension
+            vid_dim: int, video feature input dimension
+            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
+                         CG-DETR can detect in a single video.
+            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
+            contrastive_align_loss: If true, perform span - tokens contrastive learning
+            contrastive_hdim: dimension used for projecting the embeddings before computing contrastive loss
+            max_v_l: int, maximum #clips in videos
+            span_loss_type: str, one of [l1, ce]
+                l1: (center-x, width) regression.
+                ce: (st_idx, ed_idx) classification.
+            # foreground_thd: float, intersection over prediction >= foreground_thd: labeled as foreground
+            # background_thd: float, intersection over prediction <= background_thd: labeled background
+        """
+        super().__init__()
+        self.args=args
+        self.num_queries = num_queries
+        self.transformer = transformer
+        self.position_embed = position_embed
+        self.txt_position_embed = txt_position_embed
+        hidden_dim = transformer.d_model
+        self.span_loss_type = span_loss_type
+        self.max_v_l = max_v_l
+        span_pred_dim = 2 if span_loss_type == "l1" else max_v_l * 2
+        self.span_embed = MLP(hidden_dim, hidden_dim, span_pred_dim, 3)
+        self.class_embed = nn.Linear(hidden_dim, 2)  # 0: background, 1: foreground
+        self.token_type_embeddings = nn.Embedding(2, hidden_dim)
+        self.token_type_embeddings.apply(init_weights)
+        self.use_txt_pos = use_txt_pos
+        self.n_input_proj = n_input_proj
+        self.query_embed = nn.Embedding(num_queries, 2)
+        relu_args = [True] * 3
+        relu_args[n_input_proj-1] = False
+        self.input_txt_proj = nn.Sequential(*[
+            LinearLayer(txt_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[0]),
+            LinearLayer(hidden_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[1]),
+            LinearLayer(hidden_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[2])
+        ][:n_input_proj])
+        self.input_vid_proj = nn.Sequential(*[
+            LinearLayer(vid_dim + aud_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[0]),
+            LinearLayer(hidden_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[1]),
+            LinearLayer(hidden_dim, hidden_dim, layer_norm=True, dropout=input_dropout, relu=relu_args[2])
+        ][:n_input_proj])
+        self.contrastive_align_loss = contrastive_align_loss
+        if contrastive_align_loss:
+            self.contrastive_align_projection_query = nn.Linear(hidden_dim, contrastive_hdim)
+            self.contrastive_align_projection_txt = nn.Linear(hidden_dim, contrastive_hdim)
+            self.contrastive_align_projection_vid = nn.Linear(hidden_dim, contrastive_hdim)
+
+        self.saliency_proj1 = nn.Linear(hidden_dim, hidden_dim)
+        self.saliency_proj2 = nn.Linear(hidden_dim, hidden_dim)
+        self.aux_loss = aux_loss
+        self.hidden_dim = hidden_dim
+        self.global_rep_token = torch.nn.Parameter(torch.randn(args.total_prompts, hidden_dim))
+        self.global_rep_pos = torch.nn.Parameter(torch.randn(1, hidden_dim))
+        self.moment_rep_token = torch.nn.Parameter(torch.randn(hidden_dim))
+        self.moment_rep_pos = torch.nn.Parameter(torch.randn(hidden_dim))
+
+        self.dummy_rep_token = torch.nn.Parameter(torch.randn(args.num_dummies, hidden_dim))
+        self.dummy_rep_pos = torch.nn.Parameter(torch.randn(args.num_dummies, hidden_dim))
+        normalize_before = False
+        self.sent_rep_token = torch.nn.Parameter(torch.randn(hidden_dim))
+        self.sent_rep_pos = torch.nn.Parameter(torch.randn(hidden_dim))
+
+        self.txt_proj_linear = LinearLayer(txt_dim, hidden_dim, layer_norm=True)
+
+        input_txt_sa_proj = TransformerEncoderLayer(hidden_dim, 8, self.args.dim_feedforward, 0.1, "prelu", normalize_before)
+        txtproj_encoder_norm = nn.LayerNorm(hidden_dim) if normalize_before else None
+        self.txtproj_encoder = TransformerEncoder(input_txt_sa_proj, args.dummy_layers, txtproj_encoder_norm)
+
+        scls_encoder_layer = TransformerEncoderLayer(hidden_dim, 8, self.args.dim_feedforward, 0.1, "prelu", normalize_before)
+        scls_encoder_norm = nn.LayerNorm(hidden_dim) if normalize_before else None
+        self.scls_encoder = TransformerEncoder(scls_encoder_layer, args.sent_layers, scls_encoder_norm)
+
+    def forward(self, src_txt, src_txt_mask, src_vid, src_vid_mask, vid=None, qid=None, src_aud=None, src_aud_mask=None, targets=None, prompt_token=None):
+        """The forward expects two tensors:
+               - src_txt: [batch_size, L_txt, D_txt]
+               - src_txt_mask: [batch_size, L_txt], containing 0 on padded pixels,
+                    will convert to 1 as padding later for transformer
+               - src_vid: [batch_size, L_vid, D_vid]
+               - src_vid_mask: [batch_size, L_vid], containing 0 on padded pixels,
+                    will convert to 1 as padding later for transformer
+
+            It returns a dict with the following elements:
+               - "pred_spans": The normalized boxes coordinates for all queries, represented as
+                               (center_x, width). These values are normalized in [0, 1],
+                               relative to the size of each individual image (disregarding possible padding).
+                               See PostProcess for information on how to retrieve the unnormalized bounding box.
+               - "aux_outputs": Optional, only returned when auxilary losses are activated. It is a list of
+                                dictionnaries containing the two above keys for each decoder layer.
+        """
+
+   
+        ## For discovering real negative samples
+        device = src_txt_mask.device
+        # import pdb; pdb.set_trace()
+        # if vid is not None: ## for demo (run_on_video/run.py)
+        #     _count = [v.count('_') for v in vid]
+        #     if self.args.dset_name == 'hl':
+        #         _position_to_cut = [find_nth(v, '_', _count[i]-1) for i, v in enumerate(vid)]
+        #         ori_vid = [v[:_position_to_cut[i]] for i, v in enumerate(vid)]
+        #     else:
+        if vid is not None:
+            ori_vid = [v for v in vid]
+
+        if src_aud is not None:
+            src_vid = torch.cat([src_vid, src_aud], dim=2)
+
+        # --------------------------------
+        src_txt_list = []
+        src_txt_mask_list = []
+        for bs in range(src_txt.shape[0]):
+            idx = int(src_txt_mask[bs].sum().item())
+            src_txt_list.append(torch.cat((src_txt[bs, :idx, :], prompt_token[bs], src_txt[bs, idx:, :]), dim=0))
+            src_txt_mask_list.append(torch.cat((src_txt_mask[bs, :idx], torch.ones(1, dtype=torch.bfloat16).to(device), src_txt_mask[bs, idx:]), dim=0))
+        
+        src_txt = torch.stack(src_txt_list, dim=0)
+        src_txt_mask = torch.stack(src_txt_mask_list, dim=0)
+        # --------------------------------
+
+        # src_txt = torch.cat((src_txt, prompt_token), dim=1)
+        # src_txt_mask = torch.cat((src_txt_mask, torch.zeros_like(prompt_token)), dim=1)
+
+        src_vid = self.input_vid_proj(src_vid)  # [bsz,vlen,770] -> [bsz,vlen,256]
+        src_txt = self.input_txt_proj(src_txt)  # [bsz,qlen,4096] -> [bsz,qlen, 256]
+
+        src_vid = src_vid + self.token_type_embeddings(torch.full_like(src_vid_mask.long(), 1))  # TODO
+        src_txt = src_txt + self.token_type_embeddings(torch.zeros_like(src_txt_mask.long()))
+
+        # 
+        pos_vid = self.position_embed(src_vid, src_vid_mask).type(torch.bfloat16)   # (bsz, L_vid, d)
+        pos_txt = self.txt_position_embed(src_txt) if self.use_txt_pos else torch.zeros_like(src_txt).type(torch.bfloat16)  # (bsz, L_txt, d)
+
+        ### insert dummy token in front of txt
+        txt_dummy = self.dummy_rep_token.reshape([1, self.args.num_dummies, self.hidden_dim]).repeat(src_txt.shape[0], 1, 1) # [bsz, 45, 256]
+        src_txt_dummy = torch.cat([txt_dummy, src_txt], dim=1)   # [bsz, L_txt+45, 256]
+        mask_txt = torch.tensor([[True] * self.args.num_dummies]).to(src_txt_mask.device).repeat(src_txt_mask.shape[0], 1) 
+        src_txt_mask_dummy = torch.cat([mask_txt, src_txt_mask], dim=1)  # [bsz, L_txt+45]
+
+        pos_dummy = self.dummy_rep_pos.reshape([1, self.args.num_dummies, self.hidden_dim]).repeat(pos_txt.shape[0], 1, 1).type(torch.bfloat16)
+        pos_txt_dummy = torch.cat([pos_dummy, pos_txt], dim=1)
+        src_txt_dummy = src_txt_dummy.permute(1, 0, 2)   # (L, batch_size, d)
+        pos_txt_dummy = pos_txt_dummy.permute(1, 0, 2)   # (L, batch_size, d)
+
+        memory = self.txtproj_encoder(src_txt_dummy, src_key_padding_mask=~(src_txt_mask_dummy.bool()), pos=pos_txt_dummy)  # (L, batch_size, d)
+        dummy_token = memory[:self.args.num_dummies].permute(1, 0, 2)
+        pos_txt_dummy = pos_txt_dummy.permute(1, 0, 2)  # (L, batch_size, d)
+
+        src_txt_dummy = torch.cat([dummy_token, src_txt], dim=1)
+        mask_txt_dummy = torch.tensor([[True]*self.args.num_dummies]).to(src_txt_mask.device).repeat(src_txt_mask.shape[0], 1)
+        src_txt_mask_dummy = torch.cat([mask_txt_dummy, src_txt_mask], dim=1)
+
+        # Input : Concat video, dummy, txt
+        src = torch.cat([src_vid, src_txt_dummy], dim=1)  # (bsz, L_vid+L_txt, d)
+        mask = torch.cat([src_vid_mask, src_txt_mask_dummy], dim=1).bool()  # (bsz, L_vid+L_txt)
+        pos = torch.cat([pos_vid, pos_txt_dummy], dim=1)
+
+        ### sentence token
+        smask_ = torch.tensor([[True]]).to(mask.device).repeat(src_txt_mask.shape[0], 1)
+        smask = torch.cat([smask_, src_txt_mask.bool()], dim=1)
+        ssrc_ = self.sent_rep_token.reshape([1, 1, self.hidden_dim]).repeat(src_txt.shape[0], 1, 1)
+        ssrc = torch.cat([ssrc_, src_txt], dim=1)
+        spos_ = self.sent_rep_pos.reshape([1, 1, self.hidden_dim]).repeat(pos_txt.shape[0], 1, 1)
+        spos = torch.cat([spos_, pos_txt], dim=1)
+        ### dummy sentence token
+        smaskd = torch.cat([smask_, mask_txt_dummy.bool()], dim=1)
+        ssrcd = torch.cat([ssrc_, dummy_token], dim=1)
+        sposd = torch.cat([spos_, pos_dummy], dim=1)
+
+        if targets is not None: # train
+            mmask_ = torch.tensor([[True]]).to(mask.device).repeat(src_vid_mask.shape[0], 1)
+            mmask = torch.cat([mmask_, src_vid_mask.bool()], dim=1)                 # [bsz, L_vid+1]
+            moment_mask_ = torch.clamp(targets["relevant_clips"], 0, 1).bool()
+            moment_mask = torch.cat([mmask_, moment_mask_], dim=1)                  # [bsz, L_vid+1]
+            # if moment_mask.shape[1] != 76:
+            #     import pdb; pdb.set_trace()
+            mmask = mmask * moment_mask
+
+            msrc_ = self.moment_rep_token.reshape([1, 1, self.hidden_dim]).repeat(src_vid.shape[0], 1, 1)
+            msrc = torch.cat([msrc_, src_vid], dim=1)
+            mpos_ = self.moment_rep_pos.reshape([1, 1, self.hidden_dim]).repeat(pos_vid.shape[0], 1, 1)
+            mpos = torch.cat([mpos_, pos_vid], dim=1)
+
+            ### for Not moment token ####
+            nmmask_ = torch.tensor([[True]]).to(mask.device).repeat(src_vid_mask.shape[0], 1)
+            nmmask = torch.cat([nmmask_, src_vid_mask.bool()], dim=1)
+            nmoment_mask_ = ~(torch.clamp(targets["relevant_clips"], 0, 1).bool())
+            nmoment_mask = torch.cat([nmmask_, nmoment_mask_], dim=1)
+            nmmask = nmmask * nmoment_mask
+
+            nmsrc_ = self.moment_rep_token.reshape([1, 1, self.hidden_dim]).repeat(src_vid.shape[0], 1, 1)
+            nmsrc = torch.cat([nmsrc_, src_vid], dim=1)
+            nmpos_ = self.moment_rep_pos.reshape([1, 1, self.hidden_dim]).repeat(pos_vid.shape[0], 1, 1)
+            nmpos = torch.cat([nmpos_, pos_vid], dim=1)
+            ###########
+        else:
+            moment_mask_ = None
+
+        # for t2vidavg sal token
+        # import pdb; pdb.set_trace()
+        vidsrc_ = torch.zeros((len(src_vid), 1, self.hidden_dim), dtype=torch.bfloat16).to(device)
+        for i in range(len(src_vid)):
+            vidsrc_[i] = src_vid[i][:src_vid_mask.sum(1)[i].long()].mean(0).clone().detach()
+
+        video_length = src_vid.shape[1]
+        if targets is not None: ## train
+            ssrc = ssrc.permute(1, 0, 2)  # (L, batch_size, d)
+            spos = spos.permute(1, 0, 2)  # (L, batch_size, d)
+            smemory = self.scls_encoder(ssrc, src_key_padding_mask=~smask, pos=spos)  # (L, batch_size, d)
+            sentence_txt, smemory_words = smemory[0], smemory[1:] # sentence_txt : (batch_size, d)
+
+            ssrcd = ssrcd.permute(1, 0, 2)  # (L, batch_size, d)
+            sposd = sposd.permute(1, 0, 2)  # (L, batch_size, d)
+            smemoryd = self.scls_encoder(ssrcd, src_key_padding_mask=~smaskd, pos=sposd)  # (L, batch_size, d)
+            sentence_dummy, smemory_words_dummy = smemoryd[0], smemoryd[1:]
+
+            txt_dummy_proj = torch.cat([smemory_words_dummy, smemory_words], dim=0)
+
+            # import pdb; pdb.set_trace()
+            # print(src.dtype)
+            hs, reference, memory, memory_global, attn_weights, memory_moment, nmmemory_moment, mmemory_frames, nmmemory_frames = self.transformer(src, ~mask, self.query_embed.weight, pos, video_length=video_length, moment_idx=targets["relevant_clips"], msrc=msrc, mpos=mpos, mmask=~mmask, nmsrc=nmsrc, nmpos=nmpos, nmmask=~nmmask,
+                                                                                                                  ctxtoken=vidsrc_, gtoken=self.global_rep_token, gpos=self.global_rep_pos, vlen=src_vid_mask.sum(1).long())
+            moment2txt_similarity = torch.matmul(mmemory_frames.permute(1, 0, 2), txt_dummy_proj.permute(1, 2, 0))
+            nmoment2txt_similarity = torch.matmul(nmmemory_frames.permute(1, 0, 2), txt_dummy_proj.permute(1, 2, 0))
+        else: ## inference
+            sentence_dummy, sentence_txt, moment2txt_similarity, nmoment2txt_similarity = None, None, None, None
+            hs, reference, memory, memory_global, attn_weights, memory_moment, nmmemory_moment, mmemory_frames, nmmemory_frames = self.transformer(src, ~mask, self.query_embed.weight, pos, video_length=video_length,
+                                                                                                                  ctxtoken=vidsrc_, gtoken=self.global_rep_token, gpos=self.global_rep_pos, vlen=src_vid_mask.sum(1).long())
+        outputs_class = self.class_embed(hs)  # (#layers, batch_size, #queries, #classes)
+        reference_before_sigmoid = inverse_sigmoid(reference)
+        tmp = self.span_embed(hs)
+        outputs_coord = tmp + reference_before_sigmoid
+        if self.span_loss_type == "l1":
+            outputs_coord = outputs_coord.sigmoid()
+        out = {'pred_logits': outputs_class[-1], 'pred_spans': outputs_coord[-1]}
+
+        txt_mem = memory[:, src_vid.shape[1]:]  # (bsz, L_txt, d)
+        vid_mem = memory[:, :src_vid.shape[1]]  # (bsz, L_vid, d)
+        if self.contrastive_align_loss:
+            proj_queries = F.normalize(self.contrastive_align_projection_query(hs), p=2, dim=-1)
+            proj_txt_mem = F.normalize(self.contrastive_align_projection_txt(txt_mem), p=2, dim=-1)
+            proj_vid_mem = F.normalize(self.contrastive_align_projection_vid(vid_mem), p=2, dim=-1)
+            out.update(dict(
+                proj_queries=proj_queries[-1],
+                proj_txt_mem=proj_txt_mem,
+                proj_vid_mem=proj_vid_mem
+            ))
+
+        if vid is not None: ## for demo (run_on_video/run.py)
+            ### Neg Pairs ###
+            neg_vid = ori_vid[1:] + ori_vid[:1]
+            
+            real_neg_mask = torch.Tensor(element_wise_list_equal(ori_vid, neg_vid)).to(src_txt_dummy.device)
+            real_neg_mask = real_neg_mask.type(torch.bfloat16)
+
+            real_neg_mask = real_neg_mask == False
+
+            # import pdb; pdb.set_trace()
+            if real_neg_mask.sum() != 0:
+
+                src_txt_dummy_neg = torch.cat([src_txt_dummy[1:], src_txt_dummy[0:1]], dim=0)
+                src_txt_mask_dummy_neg = torch.cat([src_txt_mask_dummy[1:], src_txt_mask_dummy[0:1]], dim=0)
+                src_dummy_neg = torch.cat([src_vid, src_txt_dummy_neg], dim=1)
+                mask_dummy_neg = torch.cat([src_vid_mask, src_txt_mask_dummy_neg], dim=1).bool()
+                pos_neg = pos.clone()  # since it does not use actual content
+
+                mask_dummy_neg = mask_dummy_neg[real_neg_mask]
+                src_dummy_neg = src_dummy_neg[real_neg_mask]
+                pos_neg = pos_neg[real_neg_mask]
+                src_txt_mask_dummy_neg = src_txt_mask_dummy_neg[real_neg_mask]
+                
+                # import pdb; pdb.set_trace()
+                _, _, memory_neg, memory_global_neg, attn_weights_neg, _, _, _, _ = self.transformer(src_dummy_neg, ~mask_dummy_neg, self.query_embed.weight, pos_neg, video_length=video_length,
+                                                                                               ctxtoken=vidsrc_[real_neg_mask], gtoken=self.global_rep_token, gpos=self.global_rep_pos, vlen=src_vid_mask[real_neg_mask].sum(1).long())
+                vid_mem_neg = memory_neg[:, :src_vid.shape[1]]
+                out["saliency_scores_neg"] = (torch.sum(self.saliency_proj1(vid_mem_neg) * self.saliency_proj2(memory_global_neg).unsqueeze(1), dim=-1) / np.sqrt(self.hidden_dim))
+                out["src_txt_mask_neg"] = src_txt_mask_dummy_neg
+
+                out["t2vattnvalues_neg"] = (attn_weights_neg[:, :, self.args.num_dummies:] * (src_txt_mask_dummy_neg[:, self.args.num_dummies:].unsqueeze(1).repeat(1, video_length, 1))).sum(2)
+                out["t2vattnvalues_neg"] = torch.clamp(out["t2vattnvalues_neg"], 0, 1)
+            else:
+                out["saliency_scores_neg"] = None
+                out["t2vattnvalues_neg"] = None
+            out["real_neg_mask"] = real_neg_mask
+        else:
+            out["saliency_scores_neg"] = None
+            out["t2vattnvalues_neg"] = None
+            out["real_neg_mask"] = None
+
+
+        out["saliency_scores"] = (torch.sum(self.saliency_proj1(vid_mem) * self.saliency_proj2(memory_global).unsqueeze(1), dim=-1) / np.sqrt(self.hidden_dim))
+        out["memory_moment"] = memory_moment
+        out["nmmemory_moment"] = nmmemory_moment
+
+        ## sentence token embeeded with text / dummy
+        out["sentence_txt"] = sentence_txt
+        out["sentence_dummy"] = sentence_dummy
+        out["moment2txt_similarity"] = moment2txt_similarity
+        out["nmoment2txt_similarity"] = nmoment2txt_similarity
+        out["cate_attn_weights"] = attn_weights
+        out["moment_mask"] = moment_mask_
+        out["txt_mask"] = src_txt_mask_dummy
+
+
+        out["t2vattnvalues"] = (attn_weights[:,:,self.args.num_dummies:] * (src_txt_mask.unsqueeze(1).repeat(1, video_length, 1))).sum(2) # (batch_size, L_vid, L_txt) / (batch_size, L_txt)
+        out["t2vattnvalues"] = torch.clamp(out["t2vattnvalues"], 0, 1)
+        out["dummy_tokens"] = dummy_token
+        out["global_rep_tokens"] = self.global_rep_token
+
+        # import pdb; pdb.set_trace()
+        if targets is not None:
+            out["src_vid"] = mmemory_frames.permute(1, 0, 2) * moment_mask_.unsqueeze(2) + nmmemory_frames.permute(1, 0, 2) * (~(moment_mask_.unsqueeze(2).bool())).bfloat16()
+        else:
+            out["src_vid"] = None
+
+        out["video_mask"] = src_vid_mask
+        if self.aux_loss:
+            # assert proj_queries and proj_txt_mem
+            out['aux_outputs'] = [
+                {'pred_logits': a, 'pred_spans': b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]
+            if self.contrastive_align_loss:
+                assert proj_queries is not None
+                for idx, d in enumerate(proj_queries[:-1]):
+                    out['aux_outputs'][idx].update(dict(proj_queries=d, proj_txt_mem=proj_txt_mem))
+        return out
+
+class SetCriterion(nn.Module):
+    """ This class computes the loss for DETR.
+    The process happens in two steps:
+        1) we compute hungarian assignment between ground truth boxes and the outputs of the model
+        2) we supervise each pair of matched ground-truth / prediction (supervise class and box)
+    """
+
+    def __init__(self, matcher, weight_dict, eos_coef, losses, temperature, span_loss_type, max_v_l,
+                 saliency_margin=1, use_matcher=True, args=None):
+        """ Create the criterion.
+        Parameters:
+            matcher: module able to compute a matching between targets and proposals
+            weight_dict: dict containing as key the names of the losses and as values their relative weight.
+            eos_coef: relative classification weight applied to the no-object category
+            losses: list of all the losses to be applied. See get_loss for list of available losses.
+            temperature: float, temperature for NCE loss
+            span_loss_type: str, [l1, ce]
+            max_v_l: int,
+            saliency_margin: float
+        """
+        super().__init__()
+        self.args=args
+        self.matcher = matcher
+        self.weight_dict = weight_dict
+        self.losses = losses
+        self.temperature = temperature
+        self.span_loss_type = span_loss_type
+        self.max_v_l = max_v_l
+        self.saliency_margin = saliency_margin
+
+        # foreground and background classification
+        self.foreground_label = 0
+        self.background_label = 1
+        self.eos_coef = eos_coef
+        empty_weight = torch.ones(2)
+        empty_weight[-1] = self.eos_coef  # lower weight for background (index 1, foreground index 0)
+        self.register_buffer('empty_weight', empty_weight)
+        
+        # for tvsum,
+        self.use_matcher = use_matcher
+
+        # moment sentence contrastive
+        self.criterion = torch.nn.CrossEntropyLoss()#.to(self.args.device)
+        self.l2_criterion = torch.nn.MSELoss()#.to(self.args.device)
+        self.kld_criterion = torch.nn.KLDivLoss(reduction='none')#.to(self.args.device)
+        self.bce_criterion = nn.BCELoss(reduction='none')
+
+    def loss_spans(self, outputs, targets, indices):
+        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
+           targets dicts must contain the key "spans" containing a tensor of dim [nb_tgt_spans, 2]
+           The target spans are expected in format (center_x, w), normalized by the image size.
+        """
+        assert 'pred_spans' in outputs
+        targets = targets["span_labels"]
+        idx = self._get_src_permutation_idx(indices)
+        src_spans = outputs['pred_spans'][idx]  # (#spans, max_v_l * 2)
+        tgt_spans = torch.cat([t['spans'][i] for t, (_, i) in zip(targets, indices)], dim=0)  # (#spans, 2)
+        if self.span_loss_type == "l1":
+            loss_span = F.l1_loss(src_spans, tgt_spans, reduction='none')
+            loss_giou = 1 - torch.diag(generalized_temporal_iou(span_cxw_to_xx(src_spans), span_cxw_to_xx(tgt_spans)))
+        else:  # ce
+            n_spans = src_spans.shape[0]
+            src_spans = src_spans.view(n_spans, 2, self.max_v_l).transpose(1, 2)
+            loss_span = F.cross_entropy(src_spans, tgt_spans, reduction='none')
+            loss_giou = loss_span.new_zeros([1])
+
+        losses = {}
+        losses['loss_span'] = loss_span.mean()
+        losses['loss_giou'] = loss_giou.mean()
+        return losses
+
+    def loss_labels(self, outputs, targets, indices, log=True):
+        """Classification loss (NLL)
+        targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
+        """
+        # TODO add foreground and background classifier.  use all non-matched as background.
+        assert 'pred_logits' in outputs
+        src_logits = outputs['pred_logits']  # (batch_size, #queries, #classes=2)
+        # idx is a tuple of two 1D tensors (batch_idx, src_idx), of the same length == #objects in batch
+        idx = self._get_src_permutation_idx(indices)
+        target_classes = torch.full(src_logits.shape[:2], self.background_label,
+                                    dtype=torch.int64, device=src_logits.device)  # (batch_size, #queries)
+        target_classes[idx] = self.foreground_label
+
+        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight, reduction="none")
+        losses = {'loss_label': loss_ce.mean()}
+
+        if log:
+            # TODO this should probably be a separate loss, not hacked in this one here
+            losses['class_error'] = 100 - accuracy(src_logits[idx], self.foreground_label)[0]
+        return losses
+
+    def loss_saliency(self, outputs, targets, indices, log=True):
+        """higher scores for positive clips"""
+        if "saliency_pos_labels" not in targets:
+            return {"loss_saliency": 0}
+
+        # Neg pair loss
+        if outputs["saliency_scores_neg"] is not None: ## When batch size is not 1 (negative pair exists)
+            vid_token_mask = outputs["video_mask"]
+            real_neg_mask = outputs["real_neg_mask"]
+            saliency_scores_neg = outputs["saliency_scores_neg"].clone()  # (N, L)
+            loss_neg_pair = (- torch.log(1. - torch.sigmoid(saliency_scores_neg)) * (vid_token_mask[real_neg_mask])).sum(dim=1).mean()
+
+            saliency_scores = outputs["saliency_scores"].clone()  # (N, L)
+            saliency_contrast_label = targets["saliency_all_labels"]
+
+            # real neg
+            realneg_saliency_scores = torch.cat([saliency_scores[real_neg_mask], saliency_scores_neg], dim=1)
+            realneg_saliency_contrast_label = torch.cat([saliency_contrast_label[real_neg_mask], torch.zeros_like(saliency_contrast_label)[real_neg_mask]], dim=1)
+            realneg_vid_token_mask = vid_token_mask[real_neg_mask].repeat([1, 2])
+            realneg_saliency_scores = realneg_vid_token_mask * realneg_saliency_scores + (1. - realneg_vid_token_mask) * -1e+3
+
+            tau = 0.5
+            loss_rank_contrastive = 0.
+            for rand_idx in range(1, 12):
+                drop_mask = ~(realneg_saliency_contrast_label > 100)  # no drop
+                pos_mask = (realneg_saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                if torch.sum(pos_mask) == 0:  # no positive sample
+                    continue
+                else:
+                    batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                # drop higher ranks
+                cur_saliency_scores = realneg_saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                # numerical stability
+                logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                # softmax
+                exp_logits = torch.exp(logits)
+                log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                mean_log_prob_pos = (pos_mask * log_prob * realneg_vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                loss = - mean_log_prob_pos * batch_drop_mask
+                loss_rank_contrastive = loss_rank_contrastive + loss.mean()
+            loss_rank_contrastive = loss_rank_contrastive / 12
+
+            false_neg_mask = ~(real_neg_mask)
+            if false_neg_mask.sum() != 0:
+                if false_neg_mask.sum() == 1:
+                    falseneg_saliency_scores = saliency_scores[false_neg_mask].unsqueeze(0)
+                    falseneg_saliency_contrast_label = saliency_contrast_label[false_neg_mask].unsqueeze(0)
+                    falseneg_vid_token_mask = vid_token_mask[false_neg_mask].unsqueeze(0)
+                    falseneg_saliency_scores = falseneg_vid_token_mask * falseneg_saliency_scores + (1. - falseneg_vid_token_mask) * -1e+3
+                else:
+                    falseneg_saliency_scores = saliency_scores[false_neg_mask]
+                    falseneg_saliency_contrast_label = saliency_contrast_label[false_neg_mask]
+                    falseneg_vid_token_mask = vid_token_mask[false_neg_mask]
+                    falseneg_saliency_scores = falseneg_vid_token_mask * falseneg_saliency_scores + (1. - falseneg_vid_token_mask) * -1e+3
+
+                tau = 0.5
+                falseneg_loss_rank_contrastive = 0.
+                for rand_idx in range(1, 12):
+                    drop_mask = ~(falseneg_saliency_contrast_label > 100)  # no drop
+                    pos_mask = (falseneg_saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                    if torch.sum(pos_mask) == 0:  # no positive sample
+                        continue
+                    else:
+                        batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                    # drop higher ranks
+                    cur_saliency_scores = falseneg_saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                    # numerical stability
+                    logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                    # softmax
+                    exp_logits = torch.exp(logits)
+                    log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                    mean_log_prob_pos = (pos_mask * log_prob * falseneg_vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                    loss = - mean_log_prob_pos * batch_drop_mask
+                    falseneg_loss_rank_contrastive = falseneg_loss_rank_contrastive + loss.mean()
+                falseneg_loss_rank_contrastive = falseneg_loss_rank_contrastive / 12
+                loss_rank_contrastive += falseneg_loss_rank_contrastive
+
+            saliency_scores = outputs["saliency_scores"]  # (N, L)
+            pos_indices = targets["saliency_pos_labels"]  # (N, #pairs)
+            neg_indices = targets["saliency_neg_labels"]  # (N, #pairs)
+            num_pairs = pos_indices.shape[1]  # typically 2 or 4
+            batch_indices = torch.arange(len(saliency_scores)).to(saliency_scores.device)
+            pos_scores = torch.stack(
+                [saliency_scores[batch_indices, pos_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            neg_scores = torch.stack(
+                [saliency_scores[batch_indices, neg_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            loss_saliency = torch.clamp(self.saliency_margin + neg_scores - pos_scores, min=0).sum() \
+                            / (len(pos_scores) * num_pairs) * 2  # * 2 to keep the loss the same scale
+
+            # if self.args.dset_name in ['youtube_uni']:
+            #     loss_saliency = loss_saliency + loss_rank_contrastive + loss_neg_pair * 0.
+            # else:
+            loss_saliency = loss_saliency + loss_rank_contrastive + loss_neg_pair
+                
+            ########### Saliency loss to t2v attn weights ##############
+            """higher scores for positive clips"""
+            vid_token_mask = outputs["video_mask"]
+            # Neg pair loss
+
+            if outputs["t2vattnvalues_neg"] is not None:
+                saliency_scores_neg = outputs["t2vattnvalues_neg"].clone()  # (N, L)
+                loss_neg_pair_attn = (- torch.log(1. - saliency_scores_neg) * (vid_token_mask[real_neg_mask])).sum(dim=1).mean()
+
+            saliency_scores = outputs["t2vattnvalues"].clone()  # (N, L)
+            saliency_contrast_label = targets["saliency_all_labels"]
+
+            # real neg
+            realneg_saliency_scores = torch.cat([saliency_scores[real_neg_mask], saliency_scores_neg], dim=1)
+            realneg_saliency_contrast_label = torch.cat(
+                [saliency_contrast_label[real_neg_mask], torch.zeros_like(saliency_contrast_label)[real_neg_mask]], dim=1)
+            realneg_vid_token_mask = vid_token_mask[real_neg_mask].repeat([1, 2])
+            realneg_saliency_scores = realneg_vid_token_mask * realneg_saliency_scores + (
+                        1. - realneg_vid_token_mask) * -1e+3
+
+            tau = 0.5
+            loss_rank_contrastive_attn = 0.
+            for rand_idx in range(1, 12):
+                drop_mask = ~(realneg_saliency_contrast_label > 100)  # no drop
+                pos_mask = (realneg_saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                if torch.sum(pos_mask) == 0:  # no positive sample
+                    continue
+                else:
+                    batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                # drop higher ranks
+                cur_saliency_scores = realneg_saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                # numerical stability
+                logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                # softmax
+                exp_logits = torch.exp(logits)
+                log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                mean_log_prob_pos = (pos_mask * log_prob * realneg_vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                loss = - mean_log_prob_pos * batch_drop_mask
+                loss_rank_contrastive_attn = loss_rank_contrastive_attn + loss.mean()
+            loss_rank_contrastive_attn = loss_rank_contrastive_attn / 12
+
+            false_neg_mask = ~(real_neg_mask)
+            if false_neg_mask.sum() != 0:
+                if false_neg_mask.sum() == 1:
+                    falseneg_saliency_scores = saliency_scores[false_neg_mask].unsqueeze(0)
+                    falseneg_saliency_contrast_label = saliency_contrast_label[false_neg_mask].unsqueeze(0)
+                    falseneg_vid_token_mask = vid_token_mask[false_neg_mask].unsqueeze(0)
+                    falseneg_saliency_scores = falseneg_vid_token_mask * falseneg_saliency_scores + (1. - falseneg_vid_token_mask) * -1e+3
+                else:
+                    falseneg_saliency_scores = saliency_scores[false_neg_mask]
+                    falseneg_saliency_contrast_label = saliency_contrast_label[false_neg_mask]
+                    falseneg_vid_token_mask = vid_token_mask[false_neg_mask]
+                    falseneg_saliency_scores = falseneg_vid_token_mask * falseneg_saliency_scores + (1. - falseneg_vid_token_mask) * -1e+3
+
+                tau = 0.5
+                falseneg_loss_rank_contrastive = 0.
+                for rand_idx in range(1, 12):
+                    drop_mask = ~(falseneg_saliency_contrast_label > 100)  # no drop
+                    pos_mask = (falseneg_saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                    if torch.sum(pos_mask) == 0:  # no positive sample
+                        continue
+                    else:
+                        batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                    # drop higher ranks
+                    cur_saliency_scores = falseneg_saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                    # numerical stability
+                    logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                    # softmax
+                    exp_logits = torch.exp(logits)
+                    log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                    mean_log_prob_pos = (pos_mask * log_prob * falseneg_vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                    loss = - mean_log_prob_pos * batch_drop_mask
+                    falseneg_loss_rank_contrastive = falseneg_loss_rank_contrastive + loss.mean()
+                falseneg_loss_rank_contrastive = falseneg_loss_rank_contrastive / 12
+                loss_rank_contrastive += falseneg_loss_rank_contrastive
+
+            saliency_scores = outputs["t2vattnvalues"]  # (N, L)
+            pos_indices = targets["saliency_pos_labels"]  # (N, #pairs)
+            neg_indices = targets["saliency_neg_labels"]  # (N, #pairs)
+            num_pairs = pos_indices.shape[1]  # typically 2 or 4
+            batch_indices = torch.arange(len(saliency_scores)).to(saliency_scores.device)
+            pos_scores = torch.stack(
+                [saliency_scores[batch_indices, pos_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            neg_scores = torch.stack(
+                [saliency_scores[batch_indices, neg_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            loss_saliency_attn = torch.clamp(self.saliency_margin + neg_scores - pos_scores, min=0).sum() \
+                            / (len(pos_scores) * num_pairs) * 2  # * 2 to keep the loss the same scale
+
+            saliency_binary_label = torch.clamp(targets["saliency_all_labels"], 0, 1)
+            logits = saliency_scores.reshape(-1)
+            labels_x = saliency_binary_label.reshape(-1)
+            BCEcriterion = nn.BCELoss()
+            bceloss = BCEcriterion(logits, labels_x)
+
+            # if self.args.dset_name in ['youtube_uni']:
+            #     loss_saliency_attn = loss_rank_contrastive_attn + bceloss + loss_neg_pair_attn * 0 + loss_saliency_attn
+            # else:
+            loss_saliency_attn = loss_rank_contrastive_attn + bceloss + loss_neg_pair_attn + loss_saliency_attn
+
+            loss_saliency += (loss_saliency_attn * self.args.lw_wattn)
+            
+        else: ## when batch size == 1
+            vid_token_mask = outputs["video_mask"]
+            saliency_scores = outputs["saliency_scores"].clone()  # (N, L)
+            saliency_contrast_label = targets["saliency_all_labels"]
+
+            saliency_scores = vid_token_mask * saliency_scores + (1. - vid_token_mask) * -1e+3
+
+            tau = 0.5
+            loss_rank_contrastive = 0.
+            for rand_idx in range(1, 12):
+                drop_mask = ~(saliency_contrast_label > 100)  # no drop
+                pos_mask = (saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                if torch.sum(pos_mask) == 0:  # no positive sample
+                    continue
+                else:
+                    batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                # drop higher ranks
+                cur_saliency_scores = saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                # numerical stability
+                logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                # softmax
+                exp_logits = torch.exp(logits)
+                log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                mean_log_prob_pos = (pos_mask * log_prob * vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                loss = - mean_log_prob_pos * batch_drop_mask
+                loss_rank_contrastive = loss_rank_contrastive + loss.mean()
+            loss_rank_contrastive = loss_rank_contrastive / 12
+
+            saliency_scores = outputs["saliency_scores"]  # (N, L)
+            pos_indices = targets["saliency_pos_labels"]  # (N, #pairs)
+            neg_indices = targets["saliency_neg_labels"]  # (N, #pairs)
+            num_pairs = pos_indices.shape[1]  # typically 2 or 4
+            batch_indices = torch.arange(len(saliency_scores)).to(saliency_scores.device)
+            pos_scores = torch.stack(
+                [saliency_scores[batch_indices, pos_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            neg_scores = torch.stack(
+                [saliency_scores[batch_indices, neg_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            loss_saliency = torch.clamp(self.saliency_margin + neg_scores - pos_scores, min=0).sum() \
+                            / (len(pos_scores) * num_pairs) * 2  # * 2 to keep the loss the same scale
+
+            loss_saliency = loss_saliency + loss_rank_contrastive
+            ########### Saliency loss to t2v attn weights ##############
+            """higher scores for positive clips"""
+            vid_token_mask = outputs["video_mask"]
+            saliency_scores = outputs["t2vattnvalues"].clone()  # (N, L)
+            saliency_contrast_label = targets["saliency_all_labels"]
+
+            saliency_scores = vid_token_mask * saliency_scores + (1. - vid_token_mask) * -1e+3
+
+            tau = 0.5
+            loss_rank_contrastive = 0.
+            for rand_idx in range(1, 12):
+                drop_mask = ~(saliency_contrast_label > 100)  # no drop
+                pos_mask = (saliency_contrast_label >= rand_idx)  # positive when equal or higher than rand_idx
+                if torch.sum(pos_mask) == 0:  # no positive sample
+                    continue
+                else:
+                    batch_drop_mask = torch.sum(pos_mask, dim=1) > 0  # negative sample indicator
+
+                # drop higher ranks
+                cur_saliency_scores = saliency_scores * drop_mask / tau + ~drop_mask * -1e+3
+                # numerical stability
+                logits = cur_saliency_scores - torch.max(cur_saliency_scores, dim=1, keepdim=True)[0]
+                # softmax
+                exp_logits = torch.exp(logits)
+                log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-6)
+
+                mean_log_prob_pos = (pos_mask * log_prob * vid_token_mask).sum(1) / (pos_mask.sum(1) + 1e-6)
+                loss = - mean_log_prob_pos * batch_drop_mask
+                loss_rank_contrastive = loss_rank_contrastive + loss.mean()
+            loss_rank_contrastive_attn = loss_rank_contrastive / 12
+
+            saliency_scores = outputs["t2vattnvalues"]  # (N, L)
+            pos_indices = targets["saliency_pos_labels"]  # (N, #pairs)
+            neg_indices = targets["saliency_neg_labels"]  # (N, #pairs)
+            num_pairs = pos_indices.shape[1]  # typically 2 or 4
+            batch_indices = torch.arange(len(saliency_scores)).to(saliency_scores.device)
+            pos_scores = torch.stack(
+                [saliency_scores[batch_indices, pos_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            neg_scores = torch.stack(
+                [saliency_scores[batch_indices, neg_indices[:, col_idx]] for col_idx in range(num_pairs)], dim=1)
+            loss_saliency_attn = torch.clamp(self.saliency_margin + neg_scores - pos_scores, min=0).sum() \
+                            / (len(pos_scores) * num_pairs) * 2  # * 2 to keep the loss the same scale
+            saliency_binary_label = torch.clamp(targets["saliency_all_labels"], 0, 1)
+            logits = saliency_scores.reshape(-1)
+            labels_x = saliency_binary_label.reshape(-1)
+            BCEcriterion = nn.BCELoss()
+            bceloss = BCEcriterion(logits, labels_x)
+
+            loss_saliency_attn = loss_rank_contrastive_attn + bceloss + loss_saliency_attn
+            loss_saliency += (loss_saliency_attn * self.args.lw_wattn)
+        return {"loss_saliency": loss_saliency}
+
+    def loss_contrastive_moment_sentence(self, outputs, targets, indices, log=True):
+        if outputs["memory_moment"] is not None:
+            moment_token = outputs["memory_moment"]
+            nmmemory_moment = outputs["nmmemory_moment"]
+            sentence_token = outputs["sentence_txt"].squeeze(1)
+            sentence_dummy = outputs["sentence_dummy"].squeeze(1) # b, 1, d
+
+            moment_logits = F.normalize(moment_token, dim=1)
+            nmoment_logits = F.normalize(nmmemory_moment, dim=1)
+            sentence_logits = F.normalize(sentence_token, dim=1)
+            dummy_logits = F.normalize(sentence_dummy, dim=1)
+            # import pdb; pdb.set_trace()
+
+            similarity_matrix = torch.matmul(moment_logits, sentence_logits.T) # B B
+            nsimilarity_matrix = torch.matmul(nmoment_logits, sentence_logits.T) # B B
+            similarity_matrix = torch.cat([similarity_matrix, nsimilarity_matrix], dim=1)
+            labels = torch.eye(similarity_matrix.shape[0]).to(sentence_logits.device)
+            nlabels = torch.zeros_like(nsimilarity_matrix).to(sentence_logits.device)
+            labels = torch.cat([labels, nlabels], dim=1).max(dim=1)[1]
+
+            loss_ms_align = self.criterion(similarity_matrix, labels)
+
+            dummy_similarity_matrix = torch.matmul(moment_logits, dummy_logits.T)
+            dummy_nsimilarity_matrix = torch.matmul(nmoment_logits, dummy_logits.T)
+            dummy_similarity_matrix = torch.cat([dummy_similarity_matrix, dummy_nsimilarity_matrix], dim=1)
+            dummy_labels = (~(torch.eye(similarity_matrix.shape[0]).to(sentence_logits.device).bool())).float()
+            dummy_nlabels = torch.ones_like(nsimilarity_matrix).to(sentence_logits.device)
+            dummy_labels = torch.cat([dummy_labels, dummy_nlabels], dim=1).max(dim=1)[1]
+
+            dummy_loss_ms_align = self.criterion(dummy_similarity_matrix, dummy_labels)
+            loss_ms_align += dummy_loss_ms_align
+            video_mask = outputs['video_mask']
+            src_vid = outputs['src_vid']  # [bsz, L_vid, D_vid]
+            moment_mask_ = torch.clamp(targets["relevant_clips"], 0, 1)
+
+            momtokcls_pred = torch.matmul(moment_token.unsqueeze(1), src_vid.permute(0, 2, 1))  # bsz 1 L_vid
+            momtokcls_label = moment_mask_
+            momtokcls_logit = torch.sigmoid(momtokcls_pred)
+            loss_ms_align += (self.bce_criterion(momtokcls_logit.reshape(-1), momtokcls_label.reshape(-1)) * video_mask.reshape(-1)).mean()
+
+        else:
+            loss_ms_align = 0.
+        return {"loss_ms_align": loss_ms_align}
+        #
+
+    def loss_moment2txt_sim_distill(self, outputs, targets, indices, log=True):
+        if outputs["moment2txt_similarity"] is not None:
+            moment2txt_similarity = outputs["moment2txt_similarity"]  # bsz L_clip 22
+            moment_mask = outputs["moment_mask"].int() # bsz L_clip 1
+            txt_mask = outputs["txt_mask"].unsqueeze(1).repeat(1, outputs["cate_attn_weights"].size(1), 1)  # bsz l_t
+
+            attn_weights = outputs["cate_attn_weights"] # bsz L_clip 22
+            b, L_vid, L_txt = attn_weights.size()
+            loss_distill = self.kld_criterion(
+                torch.log(attn_weights + 1e-6).reshape(b * L_vid, -1),
+                torch.softmax(moment2txt_similarity, dim=-1).clone().detach().reshape(b * L_vid, -1)).mean(1) * moment_mask.reshape(-1)
+            loss_distill = loss_distill.sum() / moment_mask.sum()
+
+        else:
+            loss_distill = 0.
+        return {"loss_distill": loss_distill}
+
+    def loss_orthogonal_dummy(self, outputs, targets, indices, log=True):
+        dummy_tokens = outputs["dummy_tokens"]  # (n_dum, dim)
+        if dummy_tokens.size(1) != 1:
+            dummy_tokens_norm = dummy_tokens / dummy_tokens.norm(dim=2)[:, :, None]
+            dummy_tokens_sim = torch.matmul(dummy_tokens_norm, dummy_tokens_norm.permute(0, 2, 1).detach())
+            for i in range(len(dummy_tokens_sim)):
+                dummy_tokens_sim[i].fill_diagonal_(0)
+            loss_dummy_ortho = dummy_tokens_sim.abs().mean()
+        else:
+            loss_dummy_ortho=0.
+        global_tokens = outputs["global_rep_tokens"]
+
+        global_tokens_norm = global_tokens / global_tokens.norm(dim=1)[:, None]
+        global_tokens_sim = torch.matmul(global_tokens_norm, global_tokens_norm.permute(1, 0).detach())
+        for i in range(len(global_tokens_sim)):
+            global_tokens_sim.fill_diagonal_(0)
+        loss_dummy_ortho += global_tokens_sim.abs().mean()
+        return {"loss_orthogonal_dummy": loss_dummy_ortho}
+
+    def loss_contrastive_align(self, outputs, targets, indices, log=True):
+        """encourage higher scores between matched query span and input text"""
+        normalized_text_embed = outputs["proj_txt_mem"]  # (bsz, #tokens, d)  text tokens
+        normalized_img_embed = outputs["proj_queries"]  # (bsz, #queries, d)
+        logits = torch.einsum(
+            "bmd,bnd->bmn", normalized_img_embed, normalized_text_embed)  # (bsz, #queries, #tokens)
+        logits = logits.sum(2) / self.temperature  # (bsz, #queries)
+        idx = self._get_src_permutation_idx(indices)
+        positive_map = torch.zeros_like(logits, dtype=torch.bool)
+        positive_map[idx] = True
+        positive_logits = logits.masked_fill(~positive_map, 0)
+
+        pos_term = positive_logits.sum(1)  # (bsz, )
+        num_pos = positive_map.sum(1)  # (bsz, )
+        neg_term = logits.logsumexp(1)  # (bsz, )
+        loss_nce = - pos_term / num_pos + neg_term  # (bsz, )
+        losses = {"loss_contrastive_align": loss_nce.mean()}
+        return losses
+
+    def loss_contrastive_align_vid_txt(self, outputs, targets, indices, log=True):
+        """encourage higher scores between matched query span and input text"""
+        normalized_text_embed = outputs["proj_txt_mem"]  # (bsz, #tokens, d)  text tokens
+        normalized_img_embed = outputs["proj_queries"]  # (bsz, #queries, d)
+        logits = torch.einsum(
+            "bmd,bnd->bmn", normalized_img_embed, normalized_text_embed)  # (bsz, #queries, #tokens)
+        logits = logits.sum(2) / self.temperature  # (bsz, #queries)
+        idx = self._get_src_permutation_idx(indices)
+        positive_map = torch.zeros_like(logits, dtype=torch.bool)
+        positive_map[idx] = True
+        positive_logits = logits.masked_fill(~positive_map, 0)
+
+        pos_term = positive_logits.sum(1)  # (bsz, )
+        num_pos = positive_map.sum(1)  # (bsz, )
+        neg_term = logits.logsumexp(1)  # (bsz, )
+        loss_nce = - pos_term / num_pos + neg_term  # (bsz, )
+        losses = {"loss_contrastive_align": loss_nce.mean()}
+        return losses
+
+    def _get_src_permutation_idx(self, indices):
+        # permute predictions following indices
+        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+        src_idx = torch.cat([src for (src, _) in indices])
+        return batch_idx, src_idx  # two 1D tensors of the same length
+
+    def _get_tgt_permutation_idx(self, indices):
+        # permute targets following indices
+        batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+        return batch_idx, tgt_idx
+
+    def get_loss(self, loss, outputs, targets, indices, **kwargs):
+        loss_map = {
+            "spans": self.loss_spans,
+            "labels": self.loss_labels,
+            "contrastive_align": self.loss_contrastive_align,
+            "saliency": self.loss_saliency,
+            "ms_align": self.loss_contrastive_moment_sentence,
+            "distill": self.loss_moment2txt_sim_distill,
+            "orthogonal_dummy":self.loss_orthogonal_dummy
+        }
+        assert loss in loss_map, f'do you really want to compute {loss} loss?'
+        return loss_map[loss](outputs, targets, indices, **kwargs)
+
+    def forward(self, outputs, targets):
+        """ This performs the loss computation.
+        Parameters:
+             outputs: dict of tensors, see the output specification of the model for the format
+             targets: list of dicts, such that len(targets) == batch_size.
+                      The expected keys in each dict depends on the losses applied, see each loss' doc
+        """
+        outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}
+
+        # Retrieve the matching between the outputs of the last layer and the targets
+        # list(tuples), each tuple is (pred_span_indices, tgt_span_indices)
+
+        # only for HL, do not use matcher
+        if self.use_matcher:
+            # import pdb; pdb.set_trace()
+            indices = self.matcher(outputs_without_aux, targets)
+            losses_target = self.losses
+        else:
+            indices = None
+            losses_target = ["saliency"]
+
+        # Compute all the requested losses
+        losses = {}
+        for loss in losses_target:
+            losses.update(self.get_loss(loss, outputs, targets, indices))
+
+        # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+        if 'aux_outputs' in outputs:
+            for i, aux_outputs in enumerate(outputs['aux_outputs']):
+                # indices = self.matcher(aux_outputs, targets)
+                if self.use_matcher:
+                    indices = self.matcher(aux_outputs, targets)
+                    losses_target = self.losses
+                else:
+                    indices = None
+                    losses_target = ["saliency", "ms_align", "distill", "orthogonal_dummy"]
+                for loss in losses_target:
+                    if "saliency" == loss:  # skip as it is only in the top layer
+                        continue
+                    if "ms_align" == loss:
+                        continue
+                    if "distill" == loss:
+                        continue
+                    if "orthogonal_dummy" == loss:
+                        continue
+                    kwargs = {}
+                    l_dict = self.get_loss(loss, aux_outputs, targets, indices, **kwargs)
+                    l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
+                    losses.update(l_dict)
+        return losses
+
+
+class MLP(nn.Module):
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+class LinearLayer(nn.Module):
+    """linear layer configurable with layer normalization, dropout, ReLU."""
+
+    def __init__(self, input_dim, output_dim, layer_norm=True, dropout=0.1, relu=True):
+        super(LinearLayer, self).__init__()
+        self.relu = relu
+        self.layer_norm = layer_norm
+        if layer_norm:
+            self.LayerNorm = nn.LayerNorm(input_dim)
+        layers = [
+            nn.Dropout(dropout),
+            nn.Linear(input_dim, output_dim)
+        ]
+        self.net = nn.Sequential(*layers)
+
+    def forward(self, x):
+        """(N, L, D)"""
+
+        if self.layer_norm:
+            x = self.LayerNorm(x)
+        x = self.net(x)
+        if self.relu:
+            x = F.relu(x, inplace=True)
+        return x  # (N, L, D)
+
+class CGDETRConfig:
+    def __init__(self, dset_name='charadesSTA', eval_split_name='val', data_ratio=1.0,
+                 results_root='results', exp_id=None, max_es_cnt=200, eval_epoch=5,
+                 grad_clip=0.1, eval_untrained=False, resume_all=False, start_epoch=None,
+                 max_q_l=-1, max_v_l=-1, clip_length=1, max_windows=5, train_path=None,
+                 eval_path=None, no_norm_vfeat=False, no_norm_tfeat=False, v_feat_dirs=None,
+                 t_feat_dir=None, v_feat_dim=770, t_feat_dim=4096, ctx_mode='video_tef',
+                 position_embedding='sine', enc_layers=3, dec_layers=3, t2v_layers=2,
+                 sent_layers=1, moment_layers=1, dummy_layers=2, dim_feedforward=1024,
+                 hidden_dim=256, input_dropout=0.5, dropout=0.1, txt_drop_ratio=0,
+                 use_txt_pos=False, nheads=8, num_queries=10, num_dummies=45,
+                 total_prompts=10, num_prompts=1, pre_norm=False, n_input_proj=2,
+                 contrastive_hdim=64, temperature=0.07, saliency_margin=0.2, aux_loss=True,
+                 span_loss_type='l1', contrastive_align_loss=False, set_cost_span=10,
+                 set_cost_giou=1, set_cost_class=4, lw_saliency=4, lw_wattn=1.0,
+                 lw_ms_align=1.0, lw_distill=1.0, span_loss_coef=10, giou_loss_coef=1,
+                 label_loss_coef=4, eos_coef=0.1, contrastive_align_loss_coef=0.02,
+                 no_sort_results=False, max_before_nms=10, max_after_nms=10,
+                 conf_thd=0.0, nms_thd=-1):
+        
+        self.dset_name = dset_name
+        self.eval_split_name = eval_split_name
+        self.data_ratio = data_ratio
+        self.results_root = results_root
+        self.exp_id = exp_id
+        self.max_es_cnt = max_es_cnt
+        self.eval_epoch = eval_epoch
+        self.grad_clip = grad_clip
+        self.eval_untrained = eval_untrained
+        self.resume_all = resume_all
+        self.start_epoch = start_epoch
+        self.max_q_l = max_q_l
+        self.max_v_l = max_v_l
+        self.clip_length = clip_length
+        self.max_windows = max_windows
+        self.train_path = train_path
+        self.eval_path = eval_path
+        self.no_norm_vfeat = no_norm_vfeat
+        self.no_norm_tfeat = no_norm_tfeat
+        self.v_feat_dirs = v_feat_dirs
+        self.t_feat_dir = t_feat_dir
+        self.v_feat_dim = v_feat_dim
+        self.t_feat_dim = t_feat_dim
+        self.ctx_mode = ctx_mode
+        self.position_embedding = position_embedding
+        self.enc_layers = enc_layers
+        self.dec_layers = dec_layers
+        self.t2v_layers = t2v_layers
+        self.sent_layers = sent_layers
+        self.moment_layers = moment_layers
+        self.dummy_layers = dummy_layers
+        self.dim_feedforward = dim_feedforward
+        self.hidden_dim = hidden_dim
+        self.input_dropout = input_dropout
+        self.dropout = dropout
+        self.txt_drop_ratio = txt_drop_ratio
+        self.use_txt_pos = use_txt_pos
+        self.nheads = nheads
+        self.num_queries = num_queries
+        self.num_dummies = num_dummies
+        self.total_prompts = total_prompts
+        self.num_prompts = num_prompts
+        self.pre_norm = pre_norm
+        self.n_input_proj = n_input_proj
+        self.contrastive_hdim = contrastive_hdim
+        self.temperature = temperature
+        self.saliency_margin = saliency_margin
+        self.aux_loss = aux_loss
+        self.span_loss_type = span_loss_type
+        self.contrastive_align_loss = contrastive_align_loss
+        self.set_cost_span = set_cost_span
+        self.set_cost_giou = set_cost_giou
+        self.set_cost_class = set_cost_class
+        self.lw_saliency = lw_saliency
+        self.lw_wattn = lw_wattn
+        self.lw_ms_align = lw_ms_align
+        self.lw_distill = lw_distill
+        self.span_loss_coef = span_loss_coef
+        self.giou_loss_coef = giou_loss_coef
+        self.label_loss_coef = label_loss_coef
+        self.eos_coef = eos_coef
+        self.contrastive_align_loss_coef = contrastive_align_loss_coef
+        self.no_sort_results = no_sort_results
+        self.max_before_nms = max_before_nms
+        self.max_after_nms = max_after_nms
+        self.conf_thd = conf_thd
+        self.nms_thd = nms_thd
+
+def build_cgdetr_model():
+    # device = torch.device(args.device)
+    # import pdb; pdb.set_trace()
+    args = CGDETRConfig()
+
+    transformer = build_transformer(args)
+    position_embedding, txt_position_embedding = build_position_encoding(args)
+
+    # if args.a_feat_dir is None:
+    model = CGDETR(
+        transformer,
+        position_embedding,
+        txt_position_embedding,
+        txt_dim=args.t_feat_dim,
+        vid_dim=args.v_feat_dim,
+        num_queries=args.num_queries,
+        input_dropout=args.input_dropout,
+        aux_loss=args.aux_loss,
+        contrastive_align_loss=args.contrastive_align_loss,
+        contrastive_hdim=args.contrastive_hdim,
+        span_loss_type=args.span_loss_type,
+        use_txt_pos=args.use_txt_pos,
+        n_input_proj=args.n_input_proj,
+        args=args
+    )
+    # else:
+    #     model = CGDETR(
+    #         transformer,
+    #         position_embedding,
+    #         txt_position_embedding,
+    #         txt_dim=args.t_feat_dim,
+    #         vid_dim=args.v_feat_dim,
+    #         aud_dim=args.a_feat_dim,
+    #         num_queries=args.num_queries,
+    #         input_dropout=args.input_dropout,
+    #         aux_loss=args.aux_loss,
+    #         contrastive_align_loss=args.contrastive_align_loss,
+    #         contrastive_hdim=args.contrastive_hdim,
+    #         span_loss_type=args.span_loss_type,
+    #         use_txt_pos=args.use_txt_pos,
+    #         n_input_proj=args.n_input_proj,
+    #         args=args
+    #     )
+
+    matcher = build_matcher(args)
+    weight_dict = {"loss_span": args.span_loss_coef,
+                   "loss_giou": args.giou_loss_coef,
+                   "loss_label": args.label_loss_coef,
+                   "loss_saliency": args.lw_saliency,
+                   "loss_ms_align": args.lw_ms_align,
+                   "loss_distill": args.lw_distill,
+                   "loss_orthogonal_dummy":args.lw_distill}
+    if args.contrastive_align_loss:
+        weight_dict["loss_contrastive_align"] = args.contrastive_align_loss_coef
+
+    if args.aux_loss:
+        aux_weight_dict = {}
+        for i in range(args.dec_layers - 1):
+            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items() if k != "loss_saliency"})
+        weight_dict.update(aux_weight_dict)
+
+    losses = ['spans', 'labels', 'saliency', 'ms_align', 'distill', 'orthogonal_dummy']
+    if args.contrastive_align_loss:
+        losses += ["contrastive_align"]
+        
+    # For highlight detection datasets
+    # use_matcher = not (args.dset_name in ['youtube_uni', 'tvsum'])
+    use_matcher = True
+        
+    criterion = SetCriterion(
+        matcher=matcher, weight_dict=weight_dict, losses=losses,
+        eos_coef=args.eos_coef, temperature=args.temperature,
+        span_loss_type=args.span_loss_type, max_v_l=args.max_v_l,
+        saliency_margin=args.saliency_margin, use_matcher=use_matcher, args=args
+    )
+    # criterion.to(device)
+    return model, criterion

third_party/cgdetr/cg_detr/position_encoding.py

@@ -0,0 +1,116 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+
+class TrainablePositionalEncoding(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
+    def __init__(self, max_position_embeddings, hidden_size, dropout=0.1):
+        super(TrainablePositionalEncoding, self).__init__()
+        self.position_embeddings = nn.Embedding(max_position_embeddings, hidden_size)
+        self.LayerNorm = nn.LayerNorm(hidden_size)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, input_feat):
+        """
+        Args:
+            input_feat: (N, L, D)
+        """
+        bsz, seq_length = input_feat.shape[:2]
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_feat.device)
+        position_ids = position_ids.unsqueeze(0).repeat(bsz, 1)  # (N, L)
+
+        position_embeddings = self.position_embeddings(position_ids)
+
+        embeddings = self.LayerNorm(input_feat + position_embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class PositionEmbeddingSine(nn.Module):
+    """
+    This is a more standard version of the position embedding, very similar to the one
+    used by the Attention is all you need paper, generalized to work on images. (To 1D sequences)
+    """
+    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
+        super().__init__()
+        self.num_pos_feats = num_pos_feats
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, x, mask):
+        """
+        Args:
+            x: torch.tensor, (batch_size, L, d)
+            mask: torch.tensor, (batch_size, L), with 1 as valid
+
+        Returns:
+
+        """
+        assert mask is not None
+        x_embed = mask.cumsum(1, dtype=torch.float32)  # (bsz, L)
+        if self.normalize:
+            eps = 1e-6
+            x_embed = x_embed / (x_embed[:, -1:] + eps) * self.scale
+
+        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
+
+        pos_x = x_embed[:, :, None] / dim_t  # (bsz, L, num_pos_feats)
+        pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)  # (bsz, L, num_pos_feats*2)
+        # import ipdb; ipdb.set_trace()
+        return pos_x  # .permute(0, 2, 1)  # (bsz, num_pos_feats*2, L)
+
+
+class PositionEmbeddingLearned(nn.Module):
+    """
+    Absolute pos embedding, learned.
+    """
+    def __init__(self, num_pos_feats=256):
+        super().__init__()
+        self.row_embed = nn.Embedding(50, num_pos_feats)
+        self.col_embed = nn.Embedding(50, num_pos_feats)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.uniform_(self.row_embed.weight)
+        nn.init.uniform_(self.col_embed.weight)
+
+    def forward(self, x, mask):
+        h, w = x.shape[-2:]
+        i = torch.arange(w, device=x.device)
+        j = torch.arange(h, device=x.device)
+        x_emb = self.col_embed(i)
+        y_emb = self.row_embed(j)
+        pos = torch.cat([
+            x_emb.unsqueeze(0).repeat(h, 1, 1),
+            y_emb.unsqueeze(1).repeat(1, w, 1),
+        ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)
+        return pos
+
+
+def build_position_encoding(args):
+    N_steps = args.hidden_dim
+    if args.position_embedding in ('v2', 'sine'):
+        # TODO find a better way of exposing other arguments
+        position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
+    # elif args.position_embedding in ('v3', 'learned'):
+    #     position_embedding = PositionEmbeddingLearned(N_steps)
+    else:
+        raise ValueError(f"not supported {args.position_embedding}")
+    if args.max_q_l == -1:
+        args.max_q_l = 100
+    txt_pos_embed = TrainablePositionalEncoding(
+            max_position_embeddings=args.max_q_l,
+            hidden_size=args.hidden_dim, dropout=args.input_dropout)
+    return position_embedding, txt_pos_embed

third_party/cgdetr/cg_detr/postprocessing_cg_detr.py

@@ -0,0 +1,95 @@
+import pprint
+import numpy as np
+import torch
+from third_party.cgdetr.utils.basic_utils import load_jsonl
+from third_party.cgdetr.standalone_eval.eval import eval_submission
+from tqdm import tqdm
+
+
+class PostProcessorDETR:
+    def __init__(self, clip_length=2, min_ts_val=0, max_ts_val=150,
+                 min_w_l=2, max_w_l=70, move_window_method="center",
+                 process_func_names=("clip_window_l", "clip_ts", "round_multiple")):
+        self.clip_length = clip_length
+        self.min_ts_val = min_ts_val
+        self.max_ts_val = max_ts_val
+        self.min_w_l = min_w_l
+        self.max_w_l = max_w_l
+        self.move_window_method = move_window_method
+        self.process_func_names = process_func_names
+        self.name2func = dict(
+            clip_ts=self.clip_min_max_timestamps,
+            round_multiple=self.round_to_multiple_clip_lengths,
+            clip_window_l=self.clip_window_lengths
+        )
+
+    def __call__(self, lines):
+        processed_lines = []
+        for line in tqdm(lines, desc=f"convert to multiples of clip_length={self.clip_length}"):
+            windows_and_scores = torch.tensor(line["pred_relevant_windows"])
+            windows = windows_and_scores[:, :2]
+            for func_name in self.process_func_names:
+                windows = self.name2func[func_name](windows)
+            line["pred_relevant_windows"] = torch.cat(
+                [windows, windows_and_scores[:, 2:3]], dim=1).tolist()
+            line["pred_relevant_windows"] = [e[:2] + [float(f"{e[2]:.4f}")] for e in line["pred_relevant_windows"]]
+            processed_lines.append(line)
+        return processed_lines
+
+    def clip_min_max_timestamps(self, windows):
+        """
+        windows: (#windows, 2)  torch.Tensor
+        ensure timestamps for all windows is within [min_val, max_val], clip is out of boundaries.
+        """
+        return torch.clamp(windows, min=self.min_ts_val, max=self.max_ts_val)
+
+    def round_to_multiple_clip_lengths(self, windows):
+        """
+        windows: (#windows, 2)  torch.Tensor
+        ensure the final window timestamps are multiples of `clip_length`
+        """
+        return torch.round(windows / self.clip_length) * self.clip_length
+
+    def clip_window_lengths(self, windows):
+        """
+        windows: (#windows, 2)  np.ndarray
+        ensure the final window duration are within [self.min_w_l, self.max_w_l]
+        """
+        window_lengths = windows[:, 1] - windows[:, 0]
+        small_rows = window_lengths < self.min_w_l
+        if torch.sum(small_rows) > 0:
+            windows = self.move_windows(
+                windows, small_rows, self.min_w_l, move_method=self.move_window_method)
+        large_rows = window_lengths > self.max_w_l
+        if torch.sum(large_rows) > 0:
+            windows = self.move_windows(
+                windows, large_rows, self.max_w_l, move_method=self.move_window_method)
+        return windows
+
+    @classmethod
+    def move_windows(cls, windows, row_selector, new_length, move_method="left"):
+        """
+        Args:
+            windows:
+            row_selector:
+            new_length:
+            move_method: str,
+                left: keep left unchanged
+                center: keep center unchanged
+                right: keep right unchanged
+
+        Returns:
+
+        """
+        # import ipdb;
+        # ipdb.set_trace()
+        if move_method == "left":
+            windows[row_selector, 1] = windows[row_selector, 0] + new_length
+        elif move_method == "right":
+            windows[row_selector, 0] = windows[row_selector, 1] - new_length
+        elif move_method == "center":
+            center = (windows[row_selector, 1] + windows[row_selector, 0]) / 2.
+            windows[row_selector, 0] = center - new_length / 2.
+            windows[row_selector, 1] = center + new_length / 2.
+        return windows
+

third_party/cgdetr/cg_detr/scripts/charades_sta/inference.sh

@@ -0,0 +1,8 @@
+ckpt_path=$1
+eval_split_name=$2
+eval_path=data/highlight_${eval_split_name}_release.jsonl
+PYTHONPATH=$PYTHONPATH:. python cg_detr/inference.py \
+--resume ${ckpt_path} \
+--eval_split_name ${eval_split_name} \
+--eval_path ${eval_path} \
+${@:3}

third_party/cgdetr/cg_detr/scripts/charades_sta/train.sh

@@ -0,0 +1,95 @@
+dset_name=charadesSTA
+ctx_mode=video_tef
+v_feat_types=intern
+t_feat_type=intern 
+results_root=results_charades
+exp_id=exp
+
+######## data paths
+train_path=data/charades_sta/charades_sta_train_tvr_format.jsonl
+eval_path=data/charades_sta/charades_sta_test_tvr_format.jsonl
+eval_split_name=val
+
+######## setup video+text features
+feat_root=/mnt/petrelfs/lizhilin/CGDETR-main/features/charades
+
+# video features
+v_feat_dim=0
+v_feat_dirs=()
+if [[ ${v_feat_types} == *"slowfast"* ]]; then
+  v_feat_dirs+=(${feat_root}/slowfast_features)
+  (( v_feat_dim += 2304 ))  # double brackets for arithmetic op, no need to use ${v_feat_dim}
+fi
+if [[ ${v_feat_types} == *"clip"* ]]; then
+  v_feat_dirs+=(${feat_root}/clip_features)
+  (( v_feat_dim += 512 ))
+fi
+if [[ ${v_feat_types} == *"intern"* ]]; then
+  v_feat_dirs+=(${feat_root}/charade_sta_internvideo2_videoclip_6b_w1s)
+  (( v_feat_dim += 768 ))
+fi
+
+# text features
+if [[ ${t_feat_type} == "clip" ]]; then
+  t_feat_dir=${feat_root}/clip_text_features/
+  t_feat_dim=512
+fi
+if [[ ${t_feat_type} == *"intern"* ]]; then
+  t_feat_dir=(${feat_root}/charade_sta_internvideo2_llama_text_feature)
+  t_feat_dim=4096
+fi
+
+#### training
+bsz=32
+eval_bsz=32
+num_dummies=45
+num_prompts=2
+total_prompts=10
+lr_drop=400
+enc_layers=3
+dec_layers=3
+t2v_layers=2
+dummy_layers=2
+moment_layers=1
+sent_layers=1
+
+PYTHONPATH=$PYTHONPATH:. \
+srun -p video5 \
+    --preempt \
+    --job-name=${JOB_NAME} \
+    --ntasks=1 \
+    --gres=gpu:1 \
+    --ntasks-per-node=1 \
+    --cpus-per-task=8 \
+    --kill-on-bad-exit=1 \
+    python cg_detr/train.py \
+    --dset_name ${dset_name} \
+    --ctx_mode ${ctx_mode} \
+    --train_path ${train_path} \
+    --eval_path ${eval_path} \
+    --eval_split_name ${eval_split_name} \
+    --v_feat_dirs ${v_feat_dirs[@]} \
+    --v_feat_dim ${v_feat_dim} \
+    --t_feat_dir ${t_feat_dir} \
+    --t_feat_dim ${t_feat_dim} \
+    --bsz ${bsz} \
+    --results_root ${results_root} \
+    --exp_id ${exp_id} \
+    --max_v_l -1 \
+    --clip_length 1 \
+    --lr 0.0002 \
+    --lr_drop ${lr_drop} \
+    --n_epoch 200 \
+    --contrastive_align_loss_coef 0.002 \
+    --lw_saliency 4 \
+    --enc_layers ${enc_layers} \
+    --dec_layers ${dec_layers} \
+    --t2v_layers ${t2v_layers} \
+    --moment_layers ${moment_layers} \
+    --dummy_layers ${dummy_layers} \
+    --sent_layers ${sent_layers} \
+    --eval_bsz ${eval_bsz} \
+    --num_dummies ${num_dummies} \
+    --num_prompts ${num_prompts} \
+    --total_prompts ${total_prompts} \
+    ${@:1}

third_party/cgdetr/cg_detr/scripts/inference.sh

@@ -0,0 +1,11 @@
+ckpt_path=$1
+eval_split_name=$2
+eval_path=data/highlight_${eval_split_name}_release.jsonl
+echo ${ckpt_path}
+echo ${eval_split_name}
+echo ${eval_path}
+PYTHONPATH=$PYTHONPATH:. python cg_detr/inference.py \
+--resume ${ckpt_path} \
+--eval_split_name ${eval_split_name} \
+--eval_path ${eval_path} \
+${@:3}

third_party/cgdetr/cg_detr/scripts/train.sh

@@ -0,0 +1,76 @@
+dset_name=hl
+ctx_mode=video_tef
+v_feat_types=intern
+t_feat_type=intern
+results_root=results_qvhighlights
+exp_id=exp
+
+######## data paths
+train_path=data/highlight_train_release.jsonl
+eval_path=data/highlight_val_release.jsonl
+eval_split_name=val
+
+######## setup video+text features
+feat_root=../features/qvhighlight
+
+# video features
+v_feat_dim=0
+v_feat_dirs=()
+if [[ ${v_feat_types} == *"slowfast"* ]]; then
+  v_feat_dirs+=(${feat_root}/slowfast_features)
+  (( v_feat_dim += 2304 ))  # double brackets for arithmetic op, no need to use ${v_feat_dim}
+fi
+if [[ ${v_feat_types} == *"clip"* ]]; then
+  v_feat_dirs+=(${feat_root}/clip_features)
+  (( v_feat_dim += 512 ))
+fi
+if [[ ${v_feat_types} == *"intern"* ]]; then
+  v_feat_dirs+=(${feat_root}/qvhighlight_internvideo2_videoclip_6b_w2s)
+  (( v_feat_dim += 768 ))
+fi
+
+# text features
+if [[ ${t_feat_type} == "clip" ]]; then
+  t_feat_dir=${feat_root}/clip_text_features/
+  t_feat_dim=512
+fi
+if [[ ${t_feat_type} == *"intern"* ]]; then
+  t_feat_dir=(${feat_root}/qvhighlight_internvideo2_llama_text_feature)
+  t_feat_dim=4096
+fi
+
+
+#### training
+bsz=32
+enc_layers=3
+dec_layers=3
+t2v_layers=2
+moment_layers=1
+dummy_layers=2
+sent_layers=1
+max_v_l=75
+max_q_l=32
+
+PYTHONPATH=$PYTHONPATH:. python cg_detr/train.py \
+--dset_name ${dset_name} \
+--ctx_mode ${ctx_mode} \
+--train_path ${train_path} \
+--eval_path ${eval_path} \
+--eval_split_name ${eval_split_name} \
+--v_feat_dirs ${v_feat_dirs[@]} \
+--v_feat_dim ${v_feat_dim} \
+--t_feat_dir ${t_feat_dir} \
+--t_feat_dim ${t_feat_dim} \
+--bsz ${bsz} \
+--lr 0.0002 \
+--results_root ${results_root} \
+--exp_id ${exp_id} \
+--enc_layers ${enc_layers} \
+--dec_layers ${dec_layers} \
+--t2v_layers ${t2v_layers} \
+--moment_layers ${moment_layers} \
+--dummy_layers ${dummy_layers} \
+--sent_layers ${sent_layers} \
+--max_v_l ${max_v_l} \
+--max_q_l ${max_q_l} \
+${@:1}

third_party/cgdetr/cg_detr/span_utils.py

@@ -0,0 +1,127 @@
+import torch
+
+
+def span_xx_to_cxw(xx_spans):
+    """
+    Args:
+        xx_spans: tensor, (#windows, 2) or (..., 2), each row is a window of format (st, ed)
+
+    Returns:
+        cxw_spans: tensor, (#windows, 2), each row is a window of format (center=(st+ed)/2, width=(ed-st))
+    >>> spans = torch.Tensor([[0, 1], [0.2, 0.4]])
+    >>> span_xx_to_cxw(spans)
+    tensor([[0.5000, 1.0000],
+        [0.3000, 0.2000]])
+    >>> spans = torch.Tensor([[[0, 1], [0.2, 0.4]]])
+    >>> span_xx_to_cxw(spans)
+    tensor([[[0.5000, 1.0000],
+         [0.3000, 0.2000]]])
+    """
+    center = xx_spans.sum(-1) * 0.5
+    width = xx_spans[..., 1] - xx_spans[..., 0]
+    return torch.stack([center, width], dim=-1)
+
+
+def span_cxw_to_xx(cxw_spans):
+    """
+    Args:
+        cxw_spans: tensor, (#windows, 2) or (..., 2), the last dim is a row denoting a window of format (center, width)
+
+    >>> spans = torch.Tensor([[0.5000, 1.0000], [0.3000, 0.2000]])
+    >>> span_cxw_to_xx(spans)
+    tensor([[0.0000, 1.0000],
+        [0.2000, 0.4000]])
+    >>> spans = torch.Tensor([[[0.5000, 1.0000], [0.3000, 0.2000]]])
+    >>> span_cxw_to_xx(spans)
+    tensor([[[0.0000, 1.0000],
+        [0.2000, 0.4000]]])
+    """
+    x1 = cxw_spans[..., 0] - 0.5 * cxw_spans[..., 1]
+    x2 = cxw_spans[..., 0] + 0.5 * cxw_spans[..., 1]
+    return torch.stack([x1, x2], dim=-1)
+
+
+def temporal_iou(spans1, spans2):
+    """
+    Args:
+        spans1: (N, 2) torch.Tensor, each row defines a span [st, ed]
+        spans2: (M, 2) torch.Tensor, ...
+
+    Returns:
+        iou: (N, M) torch.Tensor
+        union: (N, M) torch.Tensor
+    >>> test_spans1 = torch.Tensor([[0, 0.2], [0.5, 1.0]])
+    >>> test_spans2 = torch.Tensor([[0, 0.3], [0., 1.0]])
+    >>> temporal_iou(test_spans1, test_spans2)
+    (tensor([[0.6667, 0.2000],
+         [0.0000, 0.5000]]),
+     tensor([[0.3000, 1.0000],
+             [0.8000, 1.0000]]))
+    """
+    areas1 = spans1[:, 1] - spans1[:, 0]  # (N, )
+    areas2 = spans2[:, 1] - spans2[:, 0]  # (M, )
+
+    left = torch.max(spans1[:, None, 0], spans2[:, 0])  # (N, M)
+    right = torch.min(spans1[:, None, 1], spans2[:, 1])  # (N, M)
+
+    inter = (right - left).clamp(min=0)  # (N, M)
+    union = areas1[:, None] + areas2 - inter  # (N, M)
+
+    iou = inter / union
+    return iou, union
+
+
+def temporal_intersection_over_pred(gt_spans, pred_spans):
+    """ intersection over the second input spans
+    Args:
+        gt_spans: (N, 2),
+        pred_spans: (M, 2)
+
+    Returns:
+
+    """
+    left = torch.max(gt_spans[:, None, 0], pred_spans[:, 0])
+    right = torch.min(gt_spans[:, None, 1], pred_spans[:, 1])
+
+    inter = (right - left).clamp(min=0)  # (N, M)
+    inter_over_pred = inter / (pred_spans[:, 1] - pred_spans[:, 0])
+    return inter_over_pred
+
+
+def generalized_temporal_iou(spans1, spans2):
+    """
+    Generalized IoU from https://giou.stanford.edu/
+    Also reference to DETR implementation of generalized_box_iou
+    https://github.com/facebookresearch/detr/blob/master/util/box_ops.py#L40
+
+    Args:
+        spans1: (N, 2) torch.Tensor, each row defines a span in xx format [st, ed]
+        spans2: (M, 2) torch.Tensor, ...
+
+    Returns:
+        giou: (N, M) torch.Tensor
+
+    >>> test_spans1 = torch.Tensor([[0, 0.2], [0.5, 1.0]])
+    >>> test_spans2 = torch.Tensor([[0, 0.3], [0., 1.0]])
+    >>> generalized_temporal_iou(test_spans1, test_spans2)
+    tensor([[ 0.6667,  0.2000],
+        [-0.2000,  0.5000]])
+    """
+    spans1 = spans1.float()
+    spans2 = spans2.float()
+    
+    if (spans1[:, 1] < spans1[:, 0]).all():
+        torch.save({'spans1': spans1.cpu(), 'spans2': spans2.cpu()}, 'test_spans.pt')
+        spans1[:, 1] += 0.0001
+    print(spans1)
+    assert (spans1[:, 1] >= spans1[:, 0]).all()
+    assert (spans2[:, 1] >= spans2[:, 0]).all()
+    iou, union = temporal_iou(spans1, spans2)
+
+    left = torch.min(spans1[:, None, 0], spans2[:, 0])  # (N, M)
+    right = torch.max(spans1[:, None, 1], spans2[:, 1])  # (N, M)
+    enclosing_area = (right - left).clamp(min=0)  # (N, M)
+
+    return iou - (enclosing_area - union) / enclosing_area
+
+

third_party/cgdetr/cg_detr/start_end_dataset.py

@@ -0,0 +1,383 @@
+import torch
+from torch.utils.data import Dataset
+import numpy as np
+from tqdm import tqdm
+import random
+import logging
+from os.path import join, exists
+from third_party.cgdetr.utils.basic_utils import load_jsonl, l2_normalize_np_array
+from third_party.cgdetr.utils.tensor_utils import pad_sequences_1d
+from third_party.cgdetr.cg_detr.span_utils import span_xx_to_cxw
+# from torchtext import vocab
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+class StartEndDataset(Dataset):
+    Q_FEAT_TYPES = ["pooler_output", "last_hidden_state"]
+    """One line in data loaded from data_path."
+    {
+      "qid": 7803,
+      "query": "Man in gray top walks from outside to inside.",
+      "duration": 150,
+      "vid": "RoripwjYFp8_360.0_510.0",
+      "relevant_clip_ids": [13, 14, 15, 16, 17],
+      "relevant_windows": [[26, 36]]
+    }
+    """
+
+    def __init__(self, dset_name, data_path, v_feat_dirs, q_feat_dir,
+                 q_feat_type="last_hidden_state",
+                 max_q_l=32, max_v_l=75, data_ratio=1.0, ctx_mode="video",
+                 normalize_v=True, normalize_t=True, load_labels=True,
+                 clip_len=2, max_windows=5, span_loss_type="l1", txt_drop_ratio=0,
+                 dset_domain=None):
+        self.dset_name = dset_name
+        self.data_path = data_path
+        self.data_ratio = data_ratio
+        self.v_feat_dirs = v_feat_dirs \
+            if isinstance(v_feat_dirs, list) else [v_feat_dirs]
+        self.q_feat_dir = q_feat_dir
+        self.q_feat_type = q_feat_type
+        if max_v_l == -1:
+            max_v_l = 100000000
+        if max_q_l == -1:
+            max_q_l = 100
+        self.max_q_l = max_q_l
+        self.max_v_l = max_v_l
+        self.ctx_mode = ctx_mode
+        self.use_tef = "tef" in ctx_mode
+        self.use_video = "video" in ctx_mode
+        self.normalize_t = normalize_t
+        self.normalize_v = normalize_v
+        self.load_labels = load_labels
+        self.clip_len = clip_len
+        self.max_windows = max_windows  # maximum number of windows to use as labels
+        self.span_loss_type = span_loss_type
+        self.txt_drop_ratio = txt_drop_ratio
+        if "val" in data_path or "test" in data_path:
+            assert txt_drop_ratio == 0
+
+        if self.dset_name == 'hl':
+            self.max_q_l = 32
+            self.max_v_l = 75
+            self.clip_len = 2
+
+        # checks
+        assert q_feat_type in self.Q_FEAT_TYPES
+
+        # data
+        self.data = self.load_data()
+        
+        self.use_glove = False
+        self.use_glove = 'vgg' in self.v_feat_dirs[0]
+
+        # if self.dset_name == 'charadesSTA' and self.use_glove:
+        #     self.vocab = vocab.pretrained_aliases['glove.6B.300d']()
+        #     self.vocab.itos.extend(['<unk>'])
+        #     self.vocab.stoi['<unk>'] = self.vocab.vectors.shape[0]
+        #     self.vocab.vectors = torch.cat(
+        #         (self.vocab.vectors, torch.zeros(1, self.vocab.dim)), dim=0)
+        #     self.embedding = nn.Embedding.from_pretrained(self.vocab.vectors)
+        
+
+    def load_data(self):
+        datalist = load_jsonl(self.data_path)
+        if self.data_ratio != 1:
+            n_examples = int(len(datalist) * self.data_ratio)
+            datalist = datalist[:n_examples]
+            logger.info("Using {}% of the data: {} examples"
+                        .format(self.data_ratio * 100, n_examples))
+        return datalist
+
+    def __len__(self):
+        return len(self.data)
+
+
+    def __getitem__(self, index):
+        meta = self.data[index]
+
+        model_inputs = dict()
+
+        if self.use_glove:  # False
+            model_inputs["query_feat"] = self.get_query(meta["query"])
+        else:
+            model_inputs["query_feat"] = self._get_query_feat_by_qid(meta["qid"])  # (Dq, ) or (Lq, Dq)   # [16, 4096]
+            
+
+        if self.use_video : # True
+            model_inputs["video_feat"] = self._get_video_feat_by_vid(meta["vid"])  # (Lv, Dv)
+            ctx_l = len(model_inputs["video_feat"])
+        else:
+            ctx_l = self.max_v_l
+
+
+        if self.use_tef:
+            tef_st = torch.arange(0, ctx_l, 1.0) / ctx_l
+            tef_ed = tef_st + 1.0 / ctx_l
+            tef = torch.stack([tef_st, tef_ed], dim=1)  # (Lv, 2)
+            if self.use_video :
+                model_inputs["video_feat"] = torch.cat(
+                    [model_inputs["video_feat"], tef], dim=1)  # (Lv, Dv+2)
+            else:
+                model_inputs["video_feat"] = tef
+
+
+
+        if "relevant_windows" in meta: ## For Qvhighlights test set
+            model_inputs["span_labels"] = self.get_span_labels(meta["relevant_windows"], ctx_l)  # (#windows, 2)
+            if self.dset_name in ['charadesSTA', 'tacos', 'activitynet']: ## charades, tacos, nlq
+                model_inputs["saliency_pos_labels"], model_inputs["saliency_neg_labels"], model_inputs["saliency_all_labels"] = \
+                    self.get_saliency_labels_sub_as_query(meta["relevant_windows"][0], meta["duration"], ctx_l)  # only one gt
+            elif "subs_train" not in self.data_path:
+                model_inputs["saliency_pos_labels"], model_inputs["saliency_neg_labels"], model_inputs["saliency_all_labels"] = \
+                    self.get_saliency_labels_all(meta["relevant_clip_ids"], meta["saliency_scores"], ctx_l)
+            else:
+                model_inputs["saliency_pos_labels"], model_inputs["saliency_neg_labels"], model_inputs["saliency_all_labels"] = \
+                    self.get_saliency_labels_sub_as_query(meta["relevant_windows"][0], meta["duration"], ctx_l)  # only one gt
+
+        if 'qvhighlight' or 'qvhl' in self.data_path:
+            model_inputs["relevant_clip_ids"] = meta["relevant_clip_ids"]
+        model_inputs["vid"] = meta["vid"]
+        model_inputs["qid"] = meta["qid"]
+        return dict(meta=meta, model_inputs=model_inputs)
+
+    # def get_query(self, query):
+    #     word_inds = torch.LongTensor(
+    #         [self.vocab.stoi.get(w.lower(), 400000) for w in query.split()])
+    #     return self.embedding(word_inds)
+    def get_query(self, query):
+        print("ERROR")
+        exit()
+
+    def get_saliency_labels_sub_as_query(self, gt_window, duration, ctx_l, max_n=2):
+        clip_len = duration / ctx_l
+        gt_st = int(gt_window[0] / clip_len)
+        gt_ed = max(0, min(int(gt_window[1] / clip_len), ctx_l) - 1)
+        if gt_st > gt_ed:
+            gt_st = gt_ed
+
+        if gt_st != gt_ed:
+            pos_clip_indices = random.sample(range(gt_st, gt_ed + 1), k=max_n)  # 在GT frame idx中随机选两个
+        else:
+            if self.dset_name == 'nlq':
+                pos_clip_indices = [gt_st] * 2
+            else:
+                pos_clip_indices = [gt_st, gt_st]
+
+        neg_pool = list(range(0, gt_st)) + list(range(gt_ed+1, ctx_l)) # 非GT的frame idx
+        try:
+            neg_clip_indices = random.sample(neg_pool, k=max_n)  # 在非GT frame idx中随机选两个
+        except:
+            neg_clip_indices = pos_clip_indices
+
+        # For charades_sta
+        score_array = np.zeros(ctx_l)
+        score_array[gt_st:gt_ed + 1] = 1
+
+        return pos_clip_indices, neg_clip_indices, score_array
+        
+
+    def get_saliency_labels(self, rel_clip_ids, scores, ctx_l, max_n=1, add_easy_negative=True):
+        """Sum the scores from the three annotations, then take the two clips with the
+        maximum scores as positive, and two with the minimum scores as negative.
+        Args:
+            rel_clip_ids: list(int), list of relevant clip ids
+            scores: list([anno1_score, anno2_score, anno3_score]),
+            ctx_l: int
+            max_n: int, #clips to use as positive and negative, for easy and hard negative, respectively.
+            add_easy_negative: bool, if True, sample eay negative outside the relevant_clip_ids.
+        """
+        # indices inside rel_clip_ids
+        scores = np.array(scores)  # (#rel_clips, 3)
+        agg_scores = np.sum(scores, 1)  # (#rel_clips, )
+        sort_indices = np.argsort(agg_scores)  # increasing
+
+        # indices in the whole video
+        # the min(_, ctx_l-1) here is incorrect, but should not cause
+        # much troubles since this should be rarely used.
+        hard_pos_clip_indices = [min(rel_clip_ids[idx], ctx_l-1) for idx in sort_indices[-max_n:]]
+        hard_neg_clip_indices = [min(rel_clip_ids[idx], ctx_l-1) for idx in sort_indices[:max_n]]
+        easy_pos_clip_indices = []
+        easy_neg_clip_indices = []
+        if add_easy_negative:
+            easy_neg_pool = list(set(range(ctx_l)) - set(rel_clip_ids))
+            if len(easy_neg_pool) >= max_n:
+                easy_pos_clip_indices = random.sample(rel_clip_ids, k=max_n)
+                easy_neg_clip_indices = random.sample(easy_neg_pool, k=max_n)
+            else:  # copy the hard ones
+                easy_pos_clip_indices = hard_pos_clip_indices
+                easy_neg_clip_indices = hard_neg_clip_indices
+
+        pos_clip_indices = hard_pos_clip_indices + easy_pos_clip_indices
+        neg_clip_indices = hard_neg_clip_indices + easy_neg_clip_indices
+        return pos_clip_indices, neg_clip_indices
+
+    def get_saliency_labels_all(self, rel_clip_ids, scores, ctx_l, max_n=1, add_easy_negative=True):
+        """Sum the scores from the three annotations, then take the two clips with the
+        maximum scores as positive, and two with the minimum scores as negative.
+        Args:
+            rel_clip_ids: list(int), list of relevant clip ids
+            scores: list([anno1_score, anno2_score, anno3_score]),
+            ctx_l: int
+            max_n: int, #clips to use as positive and negative, for easy and hard negative, respectively.
+            add_easy_negative: bool, if True, sample eay negative outside the relevant_clip_ids.
+        """
+        # indices inside rel_clip_ids
+        scores = np.array(scores)  # (#rel_clips, 3)
+        agg_scores = np.sum(scores, 1)  # (#rel_clips, )
+        sort_indices = np.argsort(agg_scores)  # increasing
+
+        # score_array = [min(agg_scores[idx], ctx_l-1) for idx in range(ctx_l)]
+        score_array = np.zeros(ctx_l)
+        max_len=ctx_l
+        for idx in range(len(rel_clip_ids)):
+            if rel_clip_ids[idx] >= ctx_l:
+                max_len=max(max_len,rel_clip_ids[idx])
+                # score_array_new = np.zeros(ctx_l + 1)
+                score_array_new = np.zeros(max_len+1)
+                # score_array_new[:ctx_l] = score_array
+                score_array_new[:len(score_array)] = score_array
+                score_array = score_array_new
+            score_array[rel_clip_ids[idx]] = agg_scores[idx]
+
+        # indices in the whole video
+        # the min(_, ctx_l-1) here is incorrect, but should not cause
+        # much troubles since this should be rarely used.
+        hard_pos_clip_indices = [min(rel_clip_ids[idx], ctx_l-1) for idx in sort_indices[-max_n:]]
+        hard_neg_clip_indices = [min(rel_clip_ids[idx], ctx_l-1) for idx in sort_indices[:max_n]]
+        easy_pos_clip_indices = []
+        easy_neg_clip_indices = []
+        if add_easy_negative:
+            easy_neg_pool = list(set(range(ctx_l)) - set(rel_clip_ids))
+            if len(easy_neg_pool) >= max_n:
+                easy_pos_clip_indices = random.sample(rel_clip_ids, k=max_n)
+                easy_neg_clip_indices = random.sample(easy_neg_pool, k=max_n)
+            else:  # copy the hard ones
+                easy_pos_clip_indices = hard_pos_clip_indices
+                easy_neg_clip_indices = hard_neg_clip_indices
+
+        pos_clip_indices = hard_pos_clip_indices + easy_pos_clip_indices
+        neg_clip_indices = hard_neg_clip_indices + easy_neg_clip_indices
+        return pos_clip_indices, neg_clip_indices, score_array   
+    
+    def get_span_labels(self, windows, ctx_l):
+        """
+        windows: list([st, ed]) in seconds. E.g. [[26, 36]], corresponding st_ed clip_indices [[13, 17]] (inclusive)
+            Note a maximum of `self.max_windows` windows are used.
+        returns Tensor of shape (#windows, 2), each row is [center, width] normalized by video length
+        """
+        if len(windows) > self.max_windows:
+            random.shuffle(windows)
+            windows = windows[:self.max_windows]
+        if self.span_loss_type == "l1":
+            windows = torch.Tensor(windows) / (ctx_l * self.clip_len)  # normalized windows in xx
+            windows = span_xx_to_cxw(windows)  # normalized windows in cxw
+        elif self.span_loss_type == "ce":
+            windows = torch.Tensor([
+                [int(w[0] / self.clip_len), min(int(w[1] / self.clip_len), ctx_l) - 1]
+                for w in windows]).long()  # inclusive
+        else:
+            raise NotImplementedError
+        return windows
+
+    def _get_query_feat_by_qid(self, qid):
+        # QVhighlight dataset
+        q_feat_path = join(self.q_feat_dir, f"qid{qid}.pt")
+        # q_feat = np.load(q_feat_path)[self.q_feat_type].astype(np.float32)
+        q_feat = torch.load(q_feat_path).numpy().astype(np.float32)
+        if self.q_feat_type == "last_hidden_state":
+            q_feat = q_feat[:self.max_q_l]
+        if self.normalize_t:
+            q_feat = l2_normalize_np_array(q_feat)
+        if self.txt_drop_ratio > 0:
+            q_feat = self.random_drop_rows(q_feat)
+        return torch.from_numpy(q_feat)  # (D, ) or (Lq, D)
+
+    def random_drop_rows(self, embeddings):
+        """randomly mask num_drop rows in embeddings to be zero.
+        Args:
+            embeddings: np.ndarray (L, D)
+        """
+        num_drop_rows = round(len(embeddings) * self.txt_drop_ratio)
+        if num_drop_rows > 0:
+            row_indices = np.random.choice(
+                len(embeddings), size=num_drop_rows, replace=False)
+            embeddings[row_indices] = 0
+        return embeddings
+
+    def _get_video_feat_by_vid(self, vid):
+        v_feat_list = []
+        for _feat_dir in self.v_feat_dirs:
+            try:
+                _feat_path = join(_feat_dir, f"{vid}.pt")
+                _feat = torch.load(_feat_path)["features"][:self.max_v_l].numpy().astype(np.float32)
+            except:
+                _feat_path = join(_feat_dir, f"{vid}.pt")
+                _feat = torch.load(_feat_path)[:self.max_v_l].numpy().astype(np.float32)
+            if self.normalize_v:
+                _feat = l2_normalize_np_array(_feat)
+            v_feat_list.append(_feat)
+        # some features are slightly longer than the others
+        min_len = min([len(e) for e in v_feat_list])
+        v_feat_list = [e[:min_len] for e in v_feat_list]
+        v_feat = np.concatenate(v_feat_list, axis=1)               # (vlen=34, 768)
+        return torch.from_numpy(v_feat)  # (Lv, D)
+
+
+
+def start_end_collate(batch):
+    batch_meta = [e["meta"] for e in batch]  # seems no need to collate ?
+
+    model_inputs_keys = batch[0]["model_inputs"].keys()
+    batched_data = dict()
+    for k in model_inputs_keys:
+        if k == "span_labels":
+            batched_data[k] = [dict(spans=e["model_inputs"]["span_labels"]) for e in batch]
+            continue
+        if k in ["saliency_pos_labels", "saliency_neg_labels"]:
+            batched_data[k] = torch.LongTensor([e["model_inputs"][k] for e in batch])
+            continue
+        if k == "saliency_all_labels":
+            pad_data, mask_data = pad_sequences_1d([e["model_inputs"][k] for e in batch], dtype=np.float32, fixed_length=None)
+            batched_data[k] = torch.tensor(pad_data, dtype=torch.float32)
+            continue
+        if k == 'qid':
+            batched_data[k] = [e["model_inputs"][k] for e in batch]
+            continue
+        if k == 'vid':
+            batched_data[k] = [e["model_inputs"][k] for e in batch]
+            continue
+        batched_data[k] = pad_sequences_1d(
+            [e["model_inputs"][k] for e in batch], dtype=torch.float32, fixed_length=None)
+    return batch_meta, batched_data
+
+
+def prepare_batch_inputs(batched_model_inputs):
+    model_inputs = dict(
+        src_txt=batched_model_inputs["query_feat"][0],
+        src_txt_mask=batched_model_inputs["query_feat"][1],
+        src_vid=batched_model_inputs["video_feat"][0],
+        src_vid_mask=batched_model_inputs["video_feat"][1],
+        vid=batched_model_inputs["vid"],
+        qid=batched_model_inputs["qid"],
+    )
+    targets = {}
+
+    # import pdb; pdb.set_trace()
+
+    if "span_labels" in batched_model_inputs:
+        targets["span_labels"] = [
+            dict(spans=e["spans"])
+            for e in batched_model_inputs["span_labels"]
+        ]
+    if "saliency_pos_labels" in batched_model_inputs:
+        for name in ["saliency_pos_labels", "saliency_neg_labels"]:
+            targets[name] = batched_model_inputs[name]
+
+    if "saliency_all_labels" in batched_model_inputs:
+        targets["saliency_all_labels"] = batched_model_inputs["saliency_all_labels"]
+        targets["relevant_clips"] = batched_model_inputs["saliency_all_labels"]
+    targets = None if len(targets) == 0 else targets
+    return model_inputs, targets

third_party/cgdetr/cg_detr/text_encoder.py

@@ -0,0 +1,53 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from easydict import EasyDict as edict
+from xml.model_components import BertAttention, TrainablePositionalEncoding
+
+
+class TextEncoder(nn.Module):
+    def __init__(self, hidden_size, drop, input_drop, nheads, max_position_embeddings):
+        super().__init__()
+        self.transformer_encoder = BertAttention(edict(
+            hidden_size=hidden_size,
+            intermediate_size=hidden_size,
+            hidden_dropout_prob=drop,
+            attention_probs_dropout_prob=drop,
+            num_attention_heads=nheads,
+        ))
+        self.pos_embed = TrainablePositionalEncoding(
+            max_position_embeddings=max_position_embeddings,
+            hidden_size=hidden_size,
+            dropout=input_drop,
+        )
+        self.modular_vector_mapping = nn.Linear(hidden_size, 1, bias=False)
+
+    def forward(self, feat, mask):
+        """
+        Args:
+            feat: (N, L, D=hidden_size)
+            mask: (N, L) with 1 indicates valid
+
+        Returns:
+            (N, D)
+        """
+        feat = self.pos_embed(feat)  # (N, L, D)
+        feat = self.transformer_encoder(feat, mask.unsqueeze(1))
+        att_scores = self.modular_vector_mapping(feat)  # (N, L, 1)
+        att_scores = F.softmax(mask_logits(att_scores, mask.unsqueeze(2)), dim=1)
+        pooled_feat = torch.einsum("blm,bld->bmd", att_scores, feat)  # (N, 2 or 1, D)
+        return pooled_feat.squeeze(1)
+
+
+def mask_logits(target, mask):
+    return target * mask + (1 - mask) * (-1e10)
+
+
+def build_text_encoder(args):
+    return TextEncoder(
+        hidden_size=args.hidden_dim,
+        drop=args.dropout,
+        input_drop=args.input_dropout,
+        nheads=args.nheads,
+        max_position_embeddings=args.max_q_l
+    )

third_party/cgdetr/cg_detr/train.py

@@ -0,0 +1,283 @@
+import os
+import time
+import json
+import pprint
+import random
+import numpy as np
+from tqdm import tqdm, trange
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import torch.backends.cudnn as cudnn
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+
+# import sys
+# print(sys.path)
+# sys.path.insert(os.getcwd(),0)
+# print(sys.path)
+
+from cg_detr.config import BaseOptions
+from cg_detr.start_end_dataset import StartEndDataset, start_end_collate, prepare_batch_inputs
+from cg_detr.inference import eval_epoch, start_inference, setup_model
+from utils.basic_utils import AverageMeter, dict_to_markdown
+from utils.model_utils import count_parameters
+
+
+import logging
+logger = logging.getLogger(__name__)
+logging.basicConfig(format="%(asctime)s.%(msecs)03d:%(levelname)s:%(name)s - %(message)s",
+                    datefmt="%Y-%m-%d %H:%M:%S",
+                    level=logging.INFO)
+
+
+def set_seed(seed, use_cuda=True):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if use_cuda:
+        torch.cuda.manual_seed_all(seed)
+
+
+def train_epoch(model, criterion, train_loader, optimizer, opt, epoch_i, tb_writer):
+    logger.info(f'[Epoch {epoch_i+1}]')
+    model.train()
+    criterion.train()
+
+    # init meters
+    time_meters = defaultdict(AverageMeter)
+    loss_meters = defaultdict(AverageMeter)
+
+    num_training_examples = len(train_loader)
+    timer_dataloading = time.time()
+    for batch_idx, batch in tqdm(enumerate(train_loader),
+                                 desc="Training Iteration",
+                                 total=num_training_examples):
+        time_meters["dataloading_time"].update(time.time() - timer_dataloading)
+        timer_start = time.time()
+        model_inputs, targets = prepare_batch_inputs(batch[1], opt.device, non_blocking=opt.pin_memory)
+        time_meters["prepare_inputs_time"].update(time.time() - timer_start)
+        timer_start = time.time()
+
+        outputs = model(**model_inputs, targets=targets)
+        loss_dict = criterion(outputs, targets)
+        weight_dict = criterion.weight_dict
+        losses = sum(loss_dict[k] * weight_dict[k] for k in loss_dict.keys() if k in weight_dict)
+        time_meters["model_forward_time"].update(time.time() - timer_start)
+
+        timer_start = time.time()
+        optimizer.zero_grad()
+        losses.backward()
+        if opt.grad_clip > 0:
+            nn.utils.clip_grad_norm_(model.parameters(), opt.grad_clip)
+        optimizer.step()
+        time_meters["model_backward_time"].update(time.time() - timer_start)
+
+        loss_dict["loss_overall"] = float(losses)  # for logging only
+        for k, v in loss_dict.items():
+            loss_meters[k].update(float(v) * weight_dict[k] if k in weight_dict else float(v))
+
+        timer_dataloading = time.time()
+        if opt.debug and batch_idx == 3:
+            break
+
+    # print/add logs
+    tb_writer.add_scalar("Train/lr", float(optimizer.param_groups[0]["lr"]), epoch_i+1)
+    for k, v in loss_meters.items():
+        tb_writer.add_scalar("Train/{}".format(k), v.avg, epoch_i+1)
+
+    to_write = opt.train_log_txt_formatter.format(
+        time_str=time.strftime("%Y_%m_%d_%H_%M_%S"),
+        epoch=epoch_i+1,
+        loss_str=" ".join(["{} {:.4f}".format(k, v.avg) for k, v in loss_meters.items()]))
+    with open(opt.train_log_filepath, "a") as f:
+        f.write(to_write)
+
+    logger.info("Epoch time stats:")
+    for name, meter in time_meters.items():
+        d = {k: f"{getattr(meter, k):.4f}" for k in ["max", "min", "avg"]}
+        logger.info(f"{name} ==> {d}")
+
+
+def train(model, criterion, optimizer, lr_scheduler, train_dataset, val_dataset, opt):
+    if opt.device.type == "cuda":
+        logger.info("CUDA enabled.")
+        model.to(opt.device)
+
+    tb_writer = SummaryWriter(opt.tensorboard_log_dir)
+    tb_writer.add_text("hyperparameters", dict_to_markdown(vars(opt), max_str_len=None))
+    opt.train_log_txt_formatter = "{time_str} [Epoch] {epoch:03d} [Loss] {loss_str}\n"
+    opt.eval_log_txt_formatter = "{time_str} [Epoch] {epoch:03d} [Loss] {loss_str} [Metrics] {eval_metrics_str}\n"
+
+
+    train_loader = DataLoader(
+        train_dataset,
+        collate_fn=start_end_collate,
+        batch_size=opt.bsz,
+        num_workers=opt.num_workers,
+        shuffle=True,
+        pin_memory=opt.pin_memory
+    )
+
+    prev_best_score = 0.
+    es_cnt = 0
+    # start_epoch = 0
+    if opt.start_epoch is None:
+        start_epoch = -1 if opt.eval_untrained else 0
+    else:
+        start_epoch = opt.start_epoch
+    save_submission_filename = "latest_{}_{}_preds.jsonl".format(opt.dset_name, opt.eval_split_name)
+    for epoch_i in trange(start_epoch, opt.n_epoch, desc="Epoch"):
+        if epoch_i > -1:
+            train_epoch(model, criterion, train_loader, optimizer, opt, epoch_i, tb_writer)
+            lr_scheduler.step()
+        eval_epoch_interval = opt.eval_epoch
+        if opt.eval_path is not None and (epoch_i + 1) % eval_epoch_interval == 0:
+            with torch.no_grad():
+                metrics_no_nms, metrics_nms, eval_loss_meters, latest_file_paths = \
+                    eval_epoch(model, val_dataset, opt, save_submission_filename, epoch_i, criterion, tb_writer)
+
+            # log
+            to_write = opt.eval_log_txt_formatter.format(
+                time_str=time.strftime("%Y_%m_%d_%H_%M_%S"),
+                epoch=epoch_i,
+                loss_str=" ".join(["{} {:.4f}".format(k, v.avg) for k, v in eval_loss_meters.items()]),
+                eval_metrics_str=json.dumps(metrics_no_nms))
+
+            with open(opt.eval_log_filepath, "a") as f:
+                f.write(to_write)
+            logger.info("metrics_no_nms {}".format(pprint.pformat(metrics_no_nms["brief"], indent=4)))
+            if metrics_nms is not None:
+                logger.info("metrics_nms {}".format(pprint.pformat(metrics_nms["brief"], indent=4)))
+
+            metrics = metrics_no_nms
+            for k, v in metrics["brief"].items():
+                tb_writer.add_scalar(f"Eval/{k}", float(v), epoch_i+1)
+
+            if opt.dset_name in ['hl']:
+                stop_score = metrics["brief"]["MR-full-mAP"]
+            else:
+                stop_score = (metrics["brief"]["[email protected]"] + metrics["brief"]["[email protected]"]) / 2
+
+                
+            if stop_score > prev_best_score:
+                es_cnt = 0
+                prev_best_score = stop_score
+
+                checkpoint = {
+                    "model": model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "lr_scheduler": lr_scheduler.state_dict(),
+                    "epoch": epoch_i,
+                    "opt": opt
+                }
+                torch.save(checkpoint, opt.ckpt_filepath.replace(".ckpt", "_best.ckpt"))
+
+                best_file_paths = [e.replace("latest", "best") for e in latest_file_paths]
+                for src, tgt in zip(latest_file_paths, best_file_paths):
+                    os.renames(src, tgt)
+                logger.info("The checkpoint file has been updated.")
+            else:
+                es_cnt += 1
+                if opt.max_es_cnt != -1 and es_cnt > opt.max_es_cnt:  # early stop
+                    with open(opt.train_log_filepath, "a") as f:
+                        f.write(f"Early Stop at epoch {epoch_i}")
+                    logger.info(f"\n>>>>> Early stop at epoch {epoch_i}  {prev_best_score}\n")
+                    break
+
+            # save ckpt
+            checkpoint = {
+                "model": model.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "epoch": epoch_i,
+                "opt": opt
+            }
+            torch.save(checkpoint, opt.ckpt_filepath.replace(".ckpt", "_latest.ckpt"))
+
+        # save_interval = 10 if "subs_train" in opt.train_path else 50  # smaller for pretrain
+        # if (epoch_i + 1) % save_interval == 0 or (epoch_i + 1) % opt.lr_drop == 0:  # additional copies
+        #     checkpoint = {
+        #         "model": model.state_dict(),
+        #         "optimizer": optimizer.state_dict(),
+        #         "epoch": epoch_i,
+        #         "opt": opt
+        #     }
+        #     torch.save(checkpoint, opt.ckpt_filepath.replace(".ckpt", f"_e{epoch_i:04d}.ckpt"))
+
+        if opt.debug:
+            break
+
+    tb_writer.close()
+
+
+
+def start_training():
+    logger.info("Setup config, data and model...")
+    opt = BaseOptions().parse()
+    set_seed(opt.seed)
+    if opt.debug:  # keep the model run deterministically
+        # 'cudnn.benchmark = True' enabled auto finding the best algorithm for a specific input/net config.
+        # Enable this only when input size is fixed.
+        cudnn.benchmark = False
+        cudnn.deterministic = True
+
+
+    dataset_config = dict(
+        dset_name=opt.dset_name,
+        data_path=opt.train_path,
+        v_feat_dirs=opt.v_feat_dirs,
+        q_feat_dir=opt.t_feat_dir,
+        q_feat_type="last_hidden_state",
+        max_q_l=opt.max_q_l,
+        max_v_l=opt.max_v_l,
+        ctx_mode=opt.ctx_mode,
+        data_ratio=opt.data_ratio,
+        normalize_v=not opt.no_norm_vfeat,
+        normalize_t=not opt.no_norm_tfeat,
+        clip_len=opt.clip_length,
+        max_windows=opt.max_windows,
+        span_loss_type=opt.span_loss_type,
+        txt_drop_ratio=opt.txt_drop_ratio,
+        dset_domain=opt.dset_domain,
+    )
+    dataset_config["data_path"] = opt.train_path
+    train_dataset = StartEndDataset(**dataset_config)
+    # import pdb; pdb.set_trace()
+    # train_dataset[0]
+
+    if opt.eval_path is not None:
+        dataset_config["data_path"] = opt.eval_path
+        dataset_config["txt_drop_ratio"] = 0
+        dataset_config["q_feat_dir"] = opt.t_feat_dir.replace("sub_features", "text_features")  # for pretraining
+        # dataset_config["load_labels"] = False  # uncomment to calculate eval loss
+
+        eval_dataset = StartEndDataset(**dataset_config)
+
+    else:
+        eval_dataset = None
+
+    model, criterion, optimizer, lr_scheduler = setup_model(opt)
+    logger.info(f"Model {model}")
+    count_parameters(model)
+    logger.info("Start Training...")
+    
+    train(model, criterion, optimizer, lr_scheduler, train_dataset, eval_dataset, opt)
+    
+    return opt.ckpt_filepath.replace(".ckpt", "_best.ckpt"), opt.eval_split_name, opt.eval_path, opt.debug, opt
+
+
+if __name__ == '__main__':
+    best_ckpt_path, eval_split_name, eval_path, debug, opt = start_training()
+    if not debug:
+        input_args = ["--resume", best_ckpt_path,
+                      "--eval_split_name", eval_split_name,
+                      "--eval_path", eval_path]
+
+        import sys
+        sys.argv[1:] = input_args
+        logger.info("\n\n\nFINISHED TRAINING!!!")
+        logger.info("Evaluating model at {}".format(best_ckpt_path))
+        logger.info("Input args {}".format(sys.argv[1:]))
+        start_inference(opt)

third_party/cgdetr/cg_detr/transformer.py

@@ -0,0 +1,871 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+DETR Transformer class.
+
+Copy-paste from torch.nn.Transformer with modifications:
+    * positional encodings are passed in MHattention
+    * extra LN at the end of encoder is removed
+    * decoder returns a stack of activations from all decoding layers
+"""
+import copy
+from typing import Optional
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+import math
+import numpy as np
+from .attention import MultiheadAttention
+from .crossattention import MultiheadAttention as cateattention
+
+class MLP(nn.Module):
+    """ Very simple multi-layer perceptron (also called FFN)"""
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+def inverse_sigmoid(x, eps=1e-3):
+    x = x.clamp(min=0, max=1)
+    x1 = x.clamp(min=eps)
+    x2 = (1 - x).clamp(min=eps)
+    return torch.log(x1/x2)
+
+def gen_sineembed_for_position(pos_tensor, d_model):
+    # n_query, bs, _ = pos_tensor.size()
+    # sineembed_tensor = torch.zeros(n_query, bs, 256)
+    scale = 2 * math.pi
+    dim_t = torch.arange(d_model//2, dtype=torch.float32, device=pos_tensor.device)
+    dim_t = 10000 ** (2 * (dim_t // 2) / (d_model//2))
+    center_embed = pos_tensor[:, :, 0] * scale
+    pos_x = center_embed[:, :, None] / dim_t
+    pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
+
+    span_embed = pos_tensor[:, :, 1] * scale
+    pos_w = span_embed[:, :, None] / dim_t
+    pos_w = torch.stack((pos_w[:, :, 0::2].sin(), pos_w[:, :, 1::2].cos()), dim=3).flatten(2)
+
+    pos = torch.cat((pos_x, pos_w), dim=2)
+    return pos
+
+class Transformer(nn.Module):
+
+    def __init__(self, d_model=512, nhead=8, num_queries=2, num_encoder_layers=6,
+                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False,
+                 return_intermediate_dec=False, query_dim=2,
+                 keep_query_pos=False, query_scale_type='cond_elewise',
+                 num_patterns=0,
+                 modulate_t_attn=True,
+                 bbox_embed_diff_each_layer=False, args=None
+                 ):
+        super().__init__()
+        self.args = args
+        mcls_encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        mcls_encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.mcls_encoder = TransformerEncoder(mcls_encoder_layer, args.moment_layers, mcls_encoder_norm)
+
+        t2v_encoder_layer = T2V_TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before, self.args.num_dummies)
+        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.t2v_encoder = TransformerCATEEncoder(t2v_encoder_layer, args.t2v_layers, encoder_norm)
+
+        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before)
+        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
+        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
+
+        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
+                                                dropout, activation, normalize_before, keep_query_pos=keep_query_pos)
+        decoder_norm = nn.LayerNorm(d_model)
+        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
+                                          return_intermediate=return_intermediate_dec,
+                                          d_model=d_model, query_dim=query_dim, keep_query_pos=keep_query_pos, query_scale_type=query_scale_type,
+                                          modulate_t_attn=modulate_t_attn,
+                                          bbox_embed_diff_each_layer=bbox_embed_diff_each_layer)
+
+        self._reset_parameters()
+
+        self.d_model = d_model
+        self.nhead = nhead
+        self.dec_layers = num_decoder_layers
+        self.num_queries = num_queries
+        self.num_patterns = num_patterns
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(self, src, mask, query_embed, pos_embed, video_length=None, moment_idx=None, msrc=None, mpos=None, mmask=None,
+                nmsrc=None, nmpos=None, nmmask=None,
+                ctxtoken=None, gtoken=None, gpos=None, vlen=None):
+        """
+        Args:
+            src: (batch_size, L, d)
+            mask: (batch_size, L)
+            query_embed: (#queries, d)
+            pos_embed: (batch_size, L, d) the same as src
+            video length: feature shape
+            vlen: actual video length
+        Returns:
+        """
+        # moment token
+        device = ctxtoken.device
+        if msrc is not None:
+            msrc = msrc.permute(1, 0, 2)  # (L, batch_size, d)
+            mpos = mpos.permute(1, 0, 2)  # (L, batch_size, d)
+            mmemory = self.mcls_encoder(msrc, src_key_padding_mask=mmask, pos=mpos)  # (L, batch_size, d)
+            mmemory_moment, mmemory_frames = mmemory[0], mmemory[1:]
+        else:
+            mmemory_moment = None
+            mmemory_frames = None
+        if nmsrc is not None:
+            nmsrc = nmsrc.permute(1, 0, 2)  # (L, batch_size, d)
+            nmpos = nmpos.permute(1, 0, 2)  # (L, batch_size, d)
+            nmmemory = self.mcls_encoder(nmsrc, src_key_padding_mask=nmmask, pos=nmpos)  # (L, batch_size, d)
+            nmmemory_moment, nmmemory_frames = nmmemory[0], nmmemory[1:]
+        else:
+            nmmemory_moment = None
+            nmmemory_frames = None
+
+        # flatten NxCxHxW to HWxNxC
+        bs, l, d = src.shape
+        src = src.permute(1, 0, 2)  # (L, batch_size, d)
+        pos_embed = pos_embed.permute(1, 0, 2)   # (L, batch_size, d)
+        refpoint_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)  # (#queries, batch_size, d)
+
+        # import pdb; pdb.set_trace()
+        # print(src.dtype)
+        t2v_src, attn_weights = self.t2v_encoder(src, src_key_padding_mask=mask, pos=pos_embed, video_length=video_length)  # (L, batch_size, d)
+
+        # Saliency Token
+        ## Context
+        ctx_src_ = ctxtoken.permute(1, 0, 2) # L b d
+
+        ## Distribution Token with 10 prompt tokens
+        ### Video Clip featre - context (avg) --> Find top 10 similar tokens --> weighted sum
+        # import pdb; pdb.set_trace()
+        fr_token_sim = torch.softmax(torch.matmul(F.normalize((src[:video_length] - ctx_src_).permute(1, 0, 2), dim=2), F.normalize(gtoken, dim=1).T), dim=-1)# src : b 75 d, token : 10 x d --> b 75 10
+        ### Calculate clip importance
+        frame_importance = attn_weights[:, :, self.args.num_dummies:].sum(2).clone().detach()  # b 75
+        ### Masking empty clips
+        for i in range(len(frame_importance)):
+            frame_importance[i][vlen[i]:] *= 0.
+        ### Normalize
+        frame_importance = (frame_importance / frame_importance.sum(1).unsqueeze(1)) * frame_importance.size(1)  # b 75
+        ### Scale the similarity with importance
+        fr_token_sim = fr_token_sim * frame_importance.unsqueeze(2).repeat(1, 1, fr_token_sim.size(2))  # b 75 10
+        fr_token_sim = fr_token_sim.mean(1) # b 10
+        topk_val, topkidx = torch.topk(fr_token_sim, k=self.args.num_prompts, dim=1)
+        src_ = torch.zeros((len(fr_token_sim), self.d_model), dtype=torch.bfloat16).to(device)
+        for i in range(len(fr_token_sim)):
+            src_[i] = (topk_val[i].unsqueeze(1) * gtoken[topkidx[i]]).sum(0)
+        src_ = src_.reshape(1, src.size(1), -1)
+
+        ## Add context and distribution token
+        src_ = src_ + ctx_src_
+        pos_ = gpos.reshape([1, 1, self.d_model]).repeat(1, pos_embed.shape[1], 1)
+        mask_ = torch.tensor([[False]]).to(mask.device).repeat(mask.shape[0], 1)
+
+        # import pdb; pdb.set_trace()
+        src_, _ = self.t2v_encoder(src_, src_key_padding_mask=mask_, pos=pos_,
+                                             video_length=video_length, dummy=False)  # (L, batch_size, d)
+
+        src = torch.cat([src_, t2v_src], dim=0)
+        mask = torch.cat([mask_, mask], dim=1)
+        pos_embed = torch.cat([pos_, pos_embed], dim=0)
+
+        src = src[:video_length + 1]
+        mask = mask[:, :video_length + 1]
+        pos_embed = pos_embed[:video_length + 1]
+
+        memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)  # (L, batch_size, d)
+        memory_global, memory_local = memory[0], memory[1:]
+        memory_local += memory_global.unsqueeze(0).repeat(memory_local.size(0), 1, 1)
+        mask_local = mask[:, 1:]
+        pos_embed_local = pos_embed[1:]
+
+        tgt = torch.zeros(refpoint_embed.shape[0], bs, d).to(device)
+        tgt = tgt.type(torch.bfloat16)
+
+        # import pdb; pdb.set_trace()
+        hs, references = self.decoder(tgt, memory_local, memory_key_padding_mask=mask_local, pos=pos_embed_local, refpoints_unsigmoid=refpoint_embed)  # (#layers, #queries, batch_size, d)
+        memory_local = memory_local.transpose(0, 1)  # (batch_size, L, d)
+
+        return hs, references, memory_local, memory_global, attn_weights, mmemory_moment, nmmemory_moment, mmemory_frames, nmmemory_frames
+
+
+class TransformerCATEEncoder(nn.Module):
+    def __init__(self, encoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                dummy=True,
+                **kwargs):
+        output = src
+
+        intermediate = []
+        attn_weights = None
+        for i, layer in enumerate(self.layers):
+            output, attn_weight = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos, dummy=dummy, **kwargs)
+            if attn_weights is None:
+                attn_weights = attn_weight
+            else:
+                attn_weights = attn_weights + attn_weight
+            if self.return_intermediate:
+                intermediate.append(output)
+        attn_weights /= self.num_layers
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output, attn_weights
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(self, src,
+                mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                **kwargs):
+        output = src
+
+        intermediate = []
+
+        for layer in self.layers:
+            output = layer(output, src_mask=mask,
+                           src_key_padding_mask=src_key_padding_mask, pos=pos, **kwargs)
+            if self.return_intermediate:
+                intermediate.append(output)
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False,
+                 d_model=256, query_dim=2, keep_query_pos=False, query_scale_type='cond_elewise',
+                 modulate_t_attn=False,
+                 bbox_embed_diff_each_layer=False,
+                 ):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+        assert return_intermediate
+        self.query_dim = query_dim
+
+        assert query_scale_type in ['cond_elewise', 'cond_scalar', 'fix_elewise']
+        self.query_scale_type = query_scale_type
+        if query_scale_type == 'cond_elewise':
+            self.query_scale = MLP(d_model, d_model, d_model, 2)
+        elif query_scale_type == 'cond_scalar':
+            self.query_scale = MLP(d_model, d_model, 1, 2)
+        elif query_scale_type == 'fix_elewise':
+            self.query_scale = nn.Embedding(num_layers, d_model)
+        else:
+            raise NotImplementedError("Unknown query_scale_type: {}".format(query_scale_type))
+
+        self.ref_point_head = MLP(d_model, d_model, d_model, 2)
+
+        # self.bbox_embed = None
+        # for DAB-detr
+        if bbox_embed_diff_each_layer:
+            self.bbox_embed = nn.ModuleList([MLP(d_model, d_model, 2, 3) for i in range(num_layers)])
+        else:
+            self.bbox_embed = MLP(d_model, d_model, 2, 3)
+        # init bbox_embed
+        if bbox_embed_diff_each_layer:
+            for bbox_embed in self.bbox_embed:
+                nn.init.constant_(bbox_embed.layers[-1].weight.data, 0)
+                nn.init.constant_(bbox_embed.layers[-1].bias.data, 0)
+        else:
+            nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
+            nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)
+        self.d_model = d_model
+        self.modulate_t_attn = modulate_t_attn
+        self.bbox_embed_diff_each_layer = bbox_embed_diff_each_layer
+
+        if modulate_t_attn:
+            self.ref_anchor_head = MLP(d_model, d_model, 1, 2)
+
+        if not keep_query_pos:
+            for layer_id in range(num_layers - 1):
+                self.layers[layer_id + 1].ca_qpos_proj = None
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                refpoints_unsigmoid: Optional[Tensor] = None,  # num_queries, bs, 2
+                ):
+        output = tgt
+
+        intermediate = []
+        reference_points = refpoints_unsigmoid.sigmoid()
+        ref_points = [reference_points]
+
+        # import pdb; pdb.set_trace()
+
+        for layer_id, layer in enumerate(self.layers):
+            obj_center = reference_points[..., :self.query_dim]
+            # get sine embedding for the query vector
+            query_sine_embed = gen_sineembed_for_position(obj_center, self.d_model)
+            query_sine_embed = query_sine_embed.type(torch.bfloat16)
+
+            query_pos = self.ref_point_head(query_sine_embed)
+            # For the first decoder layer, we do not apply transformation over p_s
+            if self.query_scale_type != 'fix_elewise':
+                if layer_id == 0:
+                    pos_transformation = 1
+                else:
+                    pos_transformation = self.query_scale(output)
+            else:
+                pos_transformation = self.query_scale.weight[layer_id]
+
+            # apply transformation
+            query_sine_embed = query_sine_embed * pos_transformation
+
+            # modulated HW attentions
+            if self.modulate_t_attn:
+                reft_cond = self.ref_anchor_head(output).sigmoid()  # nq, bs, 1
+
+                query_sine_embed *= (reft_cond[..., 0] / obj_center[..., 1]).unsqueeze(-1)
+
+
+            output = layer(output, memory, tgt_mask=tgt_mask,
+                           memory_mask=memory_mask,
+                           tgt_key_padding_mask=tgt_key_padding_mask,
+                           memory_key_padding_mask=memory_key_padding_mask,
+                           pos=pos, query_pos=query_pos, query_sine_embed=query_sine_embed,
+                           is_first=(layer_id == 0))
+
+            # iter update
+            if self.bbox_embed is not None:
+                if self.bbox_embed_diff_each_layer:
+                    tmp = self.bbox_embed[layer_id](output)
+                else:
+                    tmp = self.bbox_embed(output)
+                # import ipdb; ipdb.set_trace()
+                tmp[..., :self.query_dim] += inverse_sigmoid(reference_points)
+                new_reference_points = tmp[..., :self.query_dim].sigmoid()
+                if layer_id != self.num_layers - 1:
+                    ref_points.append(new_reference_points)
+                reference_points = new_reference_points.detach()
+
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            if self.bbox_embed is not None:
+                return [
+                    torch.stack(intermediate).transpose(1, 2),
+                    torch.stack(ref_points).transpose(1, 2),
+                ]
+            else:
+                return [
+                    torch.stack(intermediate).transpose(1, 2),
+                    reference_points.unsqueeze(0).transpose(1, 2)
+                ]
+
+        return output.unsqueeze(0)
+
+
+class TransformerEncoderLayerThin(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        # self.linear1 = nn.Linear(d_model, dim_feedforward)
+        # self.dropout = nn.Dropout(dropout)
+        # self.linear2 = nn.Linear(dim_feedforward, d_model)
+        self.linear = nn.Linear(d_model, d_model)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+        # self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     src,
+                     src_mask: Optional[Tensor] = None,
+                     src_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src2 = self.linear(src2)
+        src = src + self.dropout(src2)
+        src = self.norm(src)
+        # src = src + self.dropout1(src2)
+        # src = self.norm1(src)
+        # src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        # src = src + self.dropout2(src2)
+        # src = self.norm2(src)
+        return src
+
+    def forward_pre(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        """not used"""
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+    def forward(self, src,
+                src_mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class T2V_TransformerEncoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False, num_dummies=3):
+        super().__init__()
+        self.self_attn = cateattention(d_model, nhead, dropout=dropout, num_dummies=num_dummies)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = DropPath(dropout)
+        self.dropout2 = DropPath(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+        self.nhead = nhead
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     src,
+                     src_mask: Optional[Tensor] = None,
+                     src_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     video_length=None, dummy=True):
+        assert video_length is not None
+        pos_src = self.with_pos_embed(src, pos)
+        q, k, v = pos_src[:video_length], pos_src[video_length:], src[video_length:]
+
+        qmask, kmask = src_key_padding_mask[:, :video_length].unsqueeze(2), src_key_padding_mask[:, video_length:].unsqueeze(1)
+        attn_mask = torch.matmul(qmask.float(), kmask.float()).bool().repeat(self.nhead, 1, 1)
+
+        # - key_padding_mask: :math:`(S)` or :math:`(N, S)` where N is the batch size, S is the source sequence length.
+        #   If a FloatTensor is provided, it will be directly added to the value.
+        #   If a BoolTensor is provided, the positions with the
+        #   value of ``True`` will be ignored while the position with the value of ``False`` will be unchanged.
+        # - attn_mask: 2D mask :math:`(L, S)` where L is the target sequence length, S is the source sequence length.
+        #   3D mask :math:`(N*num_heads, L, S)` where N is the batch size, L is the target sequence length,
+        #   S is the source sequence length. attn_mask ensures that position i is allowed to attend the unmasked
+        #   positions. If a BoolTensor is provided, positions with ``True``
+        #   are not allowed to attend while ``False`` values will be unchanged. If a FloatTensor
+        #   is provided, it will be added to the attention weight.
+        # print(q.shape, k.shape, v.shape, attn_mask.shape, src_key_padding_mask[:, video_length + 1:].shape)
+
+        # import pdb; pdb.set_trace()
+        src2, attn_weights = self.self_attn(q, k, v, attn_mask=attn_mask, key_padding_mask=src_key_padding_mask[:, video_length:], dummy=dummy)
+
+        src2 = src[:video_length] + self.dropout1(src2)
+        src3 = self.norm1(src2)
+        src3 = self.linear2(self.dropout(self.activation(self.linear1(src3))))
+        src2 = src2 + self.dropout2(src3)
+        src2 = self.norm2(src2)
+
+        src = torch.cat([src2, src[video_length:]])
+        return src, attn_weights
+
+    def forward_pre(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None, dummy=True):
+        pass
+
+
+    def forward(self, src,
+                src_mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None, dummy=True,
+                **kwargs):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos, dummy=dummy)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos, dummy=dummy, **kwargs)
+
+class TransformerEncoderLayer(nn.Module):
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = DropPath(dropout)
+        self.dropout2 = DropPath(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self,
+                     src,
+                     src_mask: Optional[Tensor] = None,
+                     src_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
+                              key_padding_mask=src_key_padding_mask)[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward_pre(self, src,
+                    src_mask: Optional[Tensor] = None,
+                    src_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None):
+        pass
+
+    def forward(self, src,
+                src_mask: Optional[Tensor] = None,
+                src_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False, keep_query_pos=False,
+                 rm_self_attn_decoder=False):
+        super().__init__()
+        # Decoder Self-Attention
+        if not rm_self_attn_decoder:
+            self.sa_qcontent_proj = nn.Linear(d_model, d_model)
+            self.sa_qpos_proj = nn.Linear(d_model, d_model)
+            self.sa_kcontent_proj = nn.Linear(d_model, d_model)
+            self.sa_kpos_proj = nn.Linear(d_model, d_model)
+            self.sa_v_proj = nn.Linear(d_model, d_model)
+            self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, vdim=d_model)
+
+            self.norm1 = nn.LayerNorm(d_model)
+            self.dropout1 = DropPath(dropout)
+
+        # Decoder Cross-Attention
+        self.ca_qcontent_proj = nn.Linear(d_model, d_model)
+        self.ca_qpos_proj = nn.Linear(d_model, d_model)
+        self.ca_kcontent_proj = nn.Linear(d_model, d_model)
+        self.ca_kpos_proj = nn.Linear(d_model, d_model)
+        self.ca_v_proj = nn.Linear(d_model, d_model)
+        self.ca_qpos_sine_proj = nn.Linear(d_model, d_model)
+        self.cross_attn = MultiheadAttention(d_model * 2, nhead, dropout=dropout, vdim=d_model)
+
+        self.nhead = nhead
+        self.rm_self_attn_decoder = rm_self_attn_decoder
+
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout2 = DropPath(dropout)
+        self.dropout3 = DropPath(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+        self.keep_query_pos = keep_query_pos
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None,
+                query_sine_embed=None,
+                is_first=False):
+
+        # ========== Begin of Self-Attention =============
+        if not self.rm_self_attn_decoder:
+            # Apply projections here
+            # shape: num_queries x batch_size x 256
+            q_content = self.sa_qcontent_proj(tgt)  # target is the input of the first decoder layer. zero by default.
+            q_pos = self.sa_qpos_proj(query_pos)
+            k_content = self.sa_kcontent_proj(tgt)
+            k_pos = self.sa_kpos_proj(query_pos)
+            v = self.sa_v_proj(tgt)
+
+            num_queries, bs, n_model = q_content.shape
+            hw, _, _ = k_content.shape
+
+            q = q_content + q_pos
+            k = k_content + k_pos
+
+            tgt2 = self.self_attn(q, k, value=v, attn_mask=tgt_mask,
+                                  key_padding_mask=tgt_key_padding_mask)[0]
+            # ========== End of Self-Attention =============
+
+            tgt = tgt + self.dropout1(tgt2)
+            tgt = self.norm1(tgt)
+
+        # ========== Begin of Cross-Attention =============
+        # Apply projections here
+        # shape: num_queries x batch_size x 256
+        q_content = self.ca_qcontent_proj(tgt)
+        k_content = self.ca_kcontent_proj(memory)
+        v = self.ca_v_proj(memory)
+
+        num_queries, bs, n_model = q_content.shape
+        hw, _, _ = k_content.shape
+
+        k_pos = self.ca_kpos_proj(pos)
+
+        # For the first decoder layer, we concatenate the positional embedding predicted from
+        # the object query (the positional embedding) into the original query (key) in DETR.
+        if is_first or self.keep_query_pos:
+            q_pos = self.ca_qpos_proj(query_pos)
+            q = q_content + q_pos
+            k = k_content + k_pos
+        else:
+            q = q_content
+            k = k_content
+
+        q = q.view(num_queries, bs, self.nhead, n_model // self.nhead)
+        query_sine_embed = self.ca_qpos_sine_proj(query_sine_embed)
+        query_sine_embed = query_sine_embed.view(num_queries, bs, self.nhead, n_model // self.nhead)
+        q = torch.cat([q, query_sine_embed], dim=3).view(num_queries, bs, n_model * 2)
+        k = k.view(hw, bs, self.nhead, n_model // self.nhead)
+        k_pos = k_pos.view(hw, bs, self.nhead, n_model // self.nhead)
+        k = torch.cat([k, k_pos], dim=3).view(hw, bs, n_model * 2)
+
+        tgt2 = self.cross_attn(query=q,
+                               key=k,
+                               value=v, attn_mask=memory_mask,
+                               key_padding_mask=memory_key_padding_mask)[0]
+        # ========== End of Cross-Attention =============
+
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+
+class TransformerDecoderLayerThin(nn.Module):
+    """removed intermediate layer"""
+    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
+                 activation="relu", normalize_before=False):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, d_model)
+
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        # self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = DropPath(dropout)
+        self.dropout2 = DropPath(dropout)
+
+
+        # self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(self, tgt, memory,
+                     tgt_mask: Optional[Tensor] = None,
+                     memory_mask: Optional[Tensor] = None,
+                     tgt_key_padding_mask: Optional[Tensor] = None,
+                     memory_key_padding_mask: Optional[Tensor] = None,
+                     pos: Optional[Tensor] = None,
+                     query_pos: Optional[Tensor] = None):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt2 = self.linear1(tgt2)
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        return tgt
+
+    def forward_pre(self, tgt, memory,
+                    tgt_mask: Optional[Tensor] = None,
+                    memory_mask: Optional[Tensor] = None,
+                    tgt_key_padding_mask: Optional[Tensor] = None,
+                    memory_key_padding_mask: Optional[Tensor] = None,
+                    pos: Optional[Tensor] = None,
+                    query_pos: Optional[Tensor] = None):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask,
+                              key_padding_mask=tgt_key_padding_mask)[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos),
+                                   key=self.with_pos_embed(memory, pos),
+                                   value=memory, attn_mask=memory_mask,
+                                   key_padding_mask=memory_key_padding_mask)[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(self, tgt, memory,
+                tgt_mask: Optional[Tensor] = None,
+                memory_mask: Optional[Tensor] = None,
+                tgt_key_padding_mask: Optional[Tensor] = None,
+                memory_key_padding_mask: Optional[Tensor] = None,
+                pos: Optional[Tensor] = None,
+                query_pos: Optional[Tensor] = None):
+        if self.normalize_before:
+            return self.forward_pre(tgt, memory, tgt_mask, memory_mask,
+                                    tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+        return self.forward_post(tgt, memory, tgt_mask, memory_mask,
+                                 tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)
+
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def build_transformer(args):
+    return Transformer(
+        d_model=args.hidden_dim,
+        dropout=args.dropout,
+        nhead=args.nheads,
+        dim_feedforward=args.dim_feedforward,
+        num_encoder_layers=args.enc_layers,
+        num_decoder_layers=args.dec_layers,
+        normalize_before=args.pre_norm,
+        return_intermediate_dec=True,
+        activation='prelu',
+        args=args
+    )
+
+def drop_path(x, drop_prob=0.0, training=False):
+    """
+    Stochastic Depth per sample.
+    """
+    if drop_prob == 0.0 or not training:
+        return x
+
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    mask = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    mask.floor_()
+    x = x.div(keep_prob) * mask
+
+    return x
+
+class DropPath(nn.Module):
+    """
+    Drop paths per sample (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        x = x.permute(1, 0, 2)
+        res = drop_path(x, self.drop_prob, self.training)
+        return res.permute(1, 0, 2)
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    if activation == "prelu":
+        return nn.PReLU()
+    if activation == "selu":
+        return F.selu
+    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")

third_party/cgdetr/data/LICENSE

@@ -0,0 +1,437 @@
+Attribution-NonCommercial-ShareAlike 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International Public License
+("Public License"). To the extent this Public License may be
+interpreted as a contract, You are granted the Licensed Rights in
+consideration of Your acceptance of these terms and conditions, and the
+Licensor grants You such rights in consideration of benefits the
+Licensor receives from making the Licensed Material available under
+these terms and conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. BY-NC-SA Compatible License means a license listed at
+     creativecommons.org/compatiblelicenses, approved by Creative
+     Commons as essentially the equivalent of this Public License.
+
+  d. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+
+  e. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  f. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  g. License Elements means the license attributes listed in the name
+     of a Creative Commons Public License. The License Elements of this
+     Public License are Attribution, NonCommercial, and ShareAlike.
+
+  h. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  i. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  j. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  k. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  l. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  m. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  n. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. Additional offer from the Licensor -- Adapted Material.
+               Every recipient of Adapted Material from You
+               automatically receives an offer from the Licensor to
+               exercise the Licensed Rights in the Adapted Material
+               under the conditions of the Adapter's License You apply.
+
+            c. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+  b. ShareAlike.
+
+     In addition to the conditions in Section 3(a), if You Share
+     Adapted Material You produce, the following conditions also apply.
+
+       1. The Adapter's License You apply must be a Creative Commons
+          license with the same License Elements, this version or
+          later, or a BY-NC-SA Compatible License.
+
+       2. You must include the text of, or the URI or hyperlink to, the
+          Adapter's License You apply. You may satisfy this condition
+          in any reasonable manner based on the medium, means, and
+          context in which You Share Adapted Material.
+
+       3. You may not offer or impose any additional or different terms
+          or conditions on, or apply any Effective Technological
+          Measures to, Adapted Material that restrict exercise of the
+          rights granted under the Adapter's License You apply.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material,
+     including for purposes of Section 3(b); and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

third_party/cgdetr/data/README.md

@@ -0,0 +1,24 @@
+## QVHighlights Dataset
+
+Our annotation files include 3 splits: `train`, `val` and `test`. Each file is in [JSON Line](https://jsonlines.org/) format, each row of the files can be loaded as a single `dict` in Python. Below is an example of the annotation:
+
+```
+{
+    "qid": 8737, 
+    "query": "A family is playing basketball together on a green court outside.", 
+    "duration": 126, 
+    "vid": "bP5KfdFJzC4_660.0_810.0", 
+    "relevant_windows": [[0, 16]],
+    "relevant_clip_ids": [0, 1, 2, 3, 4, 5, 6, 7], 
+    "saliency_scores": [[4, 1, 1], [4, 1, 1], [4, 2, 1], [4, 3, 2], [4, 3, 2], [4, 3, 3], [4, 3, 3], [4, 3, 2]]
+}
+```
+`qid` is a unique identifier of a `query`. This query corresponds to a video identified by its video id `vid`. The `vid` is formatted as `{youtube_id}_{start_time}_{end_time}`. Use this information, one can retrieve the YouTube video from a url `https://www.youtube.com/embed/{youtube_id}?start={start_time}&end={end_time}&version=3`. For example, the video in this example is `https://www.youtube.com/embed/bP5KfdFJzC4?start=660&end=810&version=3`. 
+`duration` is an integer indicating the duration of this video.
+`relevant_windows` is the list of windows that localize the moments, each window has two numbers, one indicates the start time of the moment, another one indicates the end time. `relevant_clip_ids` is the list of ids to the segmented 2-second clips that fall into the moments specified by `relevant_windows`, starting from 0.
+`saliency_scores` contains the saliency scores annotations, each sublist corresponds to a clip in `relevant_clip_ids`. There are 3 elements in each sublist, they are the scores from three different annotators. A score of `4` means `Very Good`, while `0` means `Very Bad`.
+
+Note that the three fields `relevant_clip_ids`, `relevant_windows` and `saliency_scores` for `test` split is not included. Please refer to [../standalone_eval/README.md](../standalone_eval/README.md) for details on evaluating predictions on `test`.
+
+In addition to the annotation files, we also provided the subtitle file for our weakly supervised ASR pre-training: [subs_train.jsonl](./subs_train.jsonl). This file is formatted similarly as our annotation files, but without the `saliency_scores` entry. This file is not needed if you do not plan to pretrain models using it.
+

third_party/cgdetr/standalone_eval/README.md

@@ -0,0 +1,54 @@
+QVHighlights Evaluation and Codalab Submission
+==================
+
+### Task Definition
+Given a video and a natural language query, our task requires a system to retrieve the most relevant moments in the video, and detect the highlightness of the clips in the video. 
+
+### Evaluation
+At project root, run
+```
+bash standalone_eval/eval_sample.sh 
+```
+This command will use [eval.py](eval.py) to evaluate the provided prediction file [sample_val_preds.jsonl](sample_val_preds.jsonl), 
+the output will be written into `sample_val_preds_metrics.json`. 
+The content in this generated file should be similar if not the same as [sample_val_preds_metrics_raw.json](sample_val_preds_metrics_raw.json) file.
+
+### Format
+
+The prediction file [sample_val_preds.jsonl](sample_val_preds.jsonl) is in [JSON Line](https://jsonlines.org/) format, each row of the files can be loaded as a single `dict` in Python. Below is an example of a single line in the prediction file:
+```
+{
+  "qid": 2579,
+  "query": "A girl and her mother cooked while talking with each other on facetime.",
+  "vid": "NUsG9BgSes0_210.0_360.0",
+  "pred_relevant_windows": [
+    [0, 70, 0.9986],
+    [78, 146, 0.4138],
+    [0, 146, 0.0444],
+    ...
+  ],  
+  "pred_saliency_scores": [-0.2452, -0.3779, -0.4746, ...]
+}
+
+```
+
+
+
+| entry | description |
+| --- | ----|
+| `qid` | `int`, unique query id |
+| `query` | `str`, natural language query, not used by the evaluation script | 
+| `vid` | `str`, unique video id | 
+| `pred_relevant_windows` | `list(list)`, moment retrieval predictions. Each sublist contains 3 elements, `[start (seconds), end (seconds), score]`| 
+| `pred_saliency_scores` | `list(float)`, highlight prediction scores. The higher the better. This list should contain a score for each of the 2-second clip in the videos, and is ordered. | 
+
+
+### Codalab Submission
+To test your model's performance on `test` split, 
+please submit both `val` and `test` predictions to our 
+[Codalab evaluation server](https://codalab.lisn.upsaclay.fr/competitions/6937). 
+The submission file should be a single `.zip ` file (no enclosing folder) 
+that contains the two prediction files 
+`hl_val_submission.jsonl` and `hl_test_submission.jsonl`, each of the `*submission.jsonl` file 
+should be formatted as instructed above. 
+

third_party/cgdetr/standalone_eval/eval.py

@@ -0,0 +1,361 @@
+import numpy as np
+from collections import OrderedDict, defaultdict
+import json
+import time
+import copy
+import multiprocessing as mp
+from src.model.cgdetr_main.standalone_eval.utils import compute_average_precision_detection, \
+    compute_temporal_iou_batch_cross, compute_temporal_iou_batch_paired, load_jsonl, get_ap
+
+
+def compute_average_precision_detection_wrapper(
+        input_triple, tiou_thresholds=np.linspace(0.5, 0.95, 10)):
+    qid, ground_truth, prediction = input_triple
+    scores = compute_average_precision_detection(
+        ground_truth, prediction, tiou_thresholds=tiou_thresholds)
+    return qid, scores
+
+
+def compute_mr_ap(submission, ground_truth, iou_thds=np.linspace(0.5, 0.95, 10),
+                  max_gt_windows=None, max_pred_windows=10, num_workers=8, chunksize=50):
+    iou_thds = [float(f"{e:.2f}") for e in iou_thds]
+    pred_qid2data = defaultdict(list)
+    for d in submission:
+        pred_windows = d["pred_relevant_windows"][:max_pred_windows] \
+            if max_pred_windows is not None else d["pred_relevant_windows"]
+        qid = d["qid"]
+        for w in pred_windows:
+            pred_qid2data[qid].append({
+                "video-id": d["qid"],  # in order to use the API
+                "t-start": w[0],
+                "t-end": w[1],
+                "score": w[2]
+            })
+
+    gt_qid2data = defaultdict(list)
+    for d in ground_truth:
+        gt_windows = d["relevant_windows"][:max_gt_windows] \
+            if max_gt_windows is not None else d["relevant_windows"]
+        qid = d["qid"]
+        for w in gt_windows:
+            gt_qid2data[qid].append({
+                "video-id": d["qid"],
+                "t-start": w[0],
+                "t-end": w[1]
+            })
+    qid2ap_list = {}
+    # start_time = time.time()
+    data_triples = [[qid, gt_qid2data[qid], pred_qid2data[qid]] for qid in pred_qid2data]
+    from functools import partial
+    compute_ap_from_triple = partial(
+        compute_average_precision_detection_wrapper, tiou_thresholds=iou_thds)
+
+    if num_workers > 1:
+        with mp.Pool(num_workers) as pool:
+            for qid, scores in pool.imap_unordered(compute_ap_from_triple, data_triples, chunksize=chunksize):
+                qid2ap_list[qid] = scores
+    else:
+        for data_triple in data_triples:
+            qid, scores = compute_ap_from_triple(data_triple)
+            qid2ap_list[qid] = scores
+
+    # print(f"compute_average_precision_detection {time.time() - start_time:.2f} seconds.")
+    ap_array = np.array(list(qid2ap_list.values()))  # (#queries, #thd)
+    ap_thds = ap_array.mean(0)  # mAP at different IoU thresholds.
+    iou_thd2ap = dict(zip([str(e) for e in iou_thds], ap_thds))
+    iou_thd2ap["average"] = np.mean(ap_thds)
+    # formatting
+    iou_thd2ap = {k: float(f"{100 * v:.2f}") for k, v in iou_thd2ap.items()}
+    return iou_thd2ap
+
+
+def compute_mr_r1(submission, ground_truth, iou_thds=np.linspace(0.3, 0.95, 14)):
+    """If a predicted segment has IoU >= iou_thd with one of the 1st GT segment, we define it positive"""
+    iou_thds = [float(f"{e:.2f}") for e in iou_thds]
+    pred_qid2window = {d["qid"]: d["pred_relevant_windows"][0][:2] for d in submission}  # :2 rm scores
+    # gt_qid2window = {d["qid"]: d["relevant_windows"][0] for d in ground_truth}
+    gt_qid2window = {}
+    for d in ground_truth:
+        cur_gt_windows = d["relevant_windows"]
+        cur_qid = d["qid"]
+        cur_max_iou_idx = 0
+        if len(cur_gt_windows) > 0:  # select the GT window that has the highest IoU
+            cur_ious = compute_temporal_iou_batch_cross(
+                np.array([pred_qid2window[cur_qid]]), np.array(d["relevant_windows"])
+            )[0]
+            cur_max_iou_idx = np.argmax(cur_ious)
+        gt_qid2window[cur_qid] = cur_gt_windows[cur_max_iou_idx]
+
+    qids = list(pred_qid2window.keys())
+    pred_windows = np.array([pred_qid2window[k] for k in qids]).astype(float)
+    gt_windows = np.array([gt_qid2window[k] for k in qids]).astype(float)
+    pred_gt_iou = compute_temporal_iou_batch_paired(pred_windows, gt_windows)
+    iou_thd2recall_at_one = {}
+    miou_at_one = float(f"{np.mean(pred_gt_iou) * 100:.2f}")
+    for thd in iou_thds:
+        iou_thd2recall_at_one[str(thd)] = float(f"{np.mean(pred_gt_iou >= thd) * 100:.2f}")
+    return iou_thd2recall_at_one, miou_at_one
+
+
+def get_window_len(window):
+    return window[1] - window[0]
+
+
+def get_data_by_range(submission, ground_truth, len_range):
+    """ keep queries with ground truth window length in the specified length range.
+    Args:
+        submission:
+        ground_truth:
+        len_range: [min_l (int), max_l (int)]. the range is (min_l, max_l], i.e., min_l < l <= max_l
+    """
+    min_l, max_l = len_range
+    if min_l == 0 and max_l == 150:  # min and max l in dataset
+        return submission, ground_truth
+
+    # only keep ground truth with windows in the specified length range
+    # if multiple GT windows exists, we only keep the ones in the range
+    ground_truth_in_range = []
+    gt_qids_in_range = set()
+    for d in ground_truth:
+        rel_windows_in_range = [
+            w for w in d["relevant_windows"] if min_l < get_window_len(w) <= max_l]
+        if len(rel_windows_in_range) > 0:
+            d = copy.deepcopy(d)
+            d["relevant_windows"] = rel_windows_in_range
+            ground_truth_in_range.append(d)
+            gt_qids_in_range.add(d["qid"])
+
+    # keep only submissions for ground_truth_in_range
+    submission_in_range = []
+    for d in submission:
+        if d["qid"] in gt_qids_in_range:
+            submission_in_range.append(copy.deepcopy(d))
+
+    return submission_in_range, ground_truth_in_range
+
+
+def eval_moment_retrieval(submission, ground_truth, verbose=True):
+    length_ranges = [[0, 10], [10, 30], [30, 150], [0, 150], ]  #
+    range_names = ["short", "middle", "long", "full"]
+
+    ret_metrics = {}
+    for l_range, name in zip(length_ranges, range_names):
+        if verbose:
+            start_time = time.time()
+        _submission, _ground_truth = get_data_by_range(submission, ground_truth, l_range)
+        print(f"{name}: {l_range}, {len(_ground_truth)}/{len(ground_truth)}="
+              f"{100*len(_ground_truth)/len(ground_truth):.2f} examples.")
+        if len(_ground_truth) == 0:
+            # ret_metrics[name] = {"MR-mAP": 0., "MR-R1": 0.}
+            dummy_dict = {}
+            for k in np.linspace(0.5, 0.95, 19):
+                dummy_dict[k] = 0.
+            dummy_dict['average'] = 0.
+            ret_metrics[name] = {"MR-mAP": dummy_dict, "MR-R1": dummy_dict}
+        else:
+            iou_thd2average_precision = compute_mr_ap(_submission, _ground_truth, num_workers=8, chunksize=50)
+            iou_thd2recall_at_one, miou_at_one = compute_mr_r1(_submission, _ground_truth)
+            ret_metrics[name] = {"MR-mIoU": miou_at_one,
+                                 "MR-mAP": iou_thd2average_precision,
+                                 "MR-R1": iou_thd2recall_at_one}
+
+            # iou_thd2average_precision = compute_mr_ap(_submission, _ground_truth, num_workers=8, chunksize=50)
+            # iou_thd2recall_at_one = compute_mr_r1(_submission, _ground_truth)
+            # ret_metrics[name] = {"MR-mAP": iou_thd2average_precision, "MR-R1": iou_thd2recall_at_one}
+            if verbose:
+                print(f"[eval_moment_retrieval] [{name}] {time.time() - start_time:.2f} seconds")
+    return ret_metrics
+
+
+def compute_hl_hit1(qid2preds, qid2gt_scores_binary):
+    qid2max_scored_clip_idx = {k: np.argmax(v["pred_saliency_scores"]) for k, v in qid2preds.items()}
+    hit_scores = np.zeros((len(qid2preds), 3))
+    qids = list(qid2preds.keys())
+    for idx, qid in enumerate(qids):
+        pred_clip_idx = qid2max_scored_clip_idx[qid]
+        gt_scores_binary = qid2gt_scores_binary[qid]   # (#clips, 3)
+        if pred_clip_idx < len(gt_scores_binary):
+            hit_scores[idx] = gt_scores_binary[pred_clip_idx]
+    # aggregate scores from 3 separate annotations (3 workers) by taking the max.
+    # then average scores from all queries.
+    hit_at_one = float(f"{100 * np.mean(np.max(hit_scores, 1)):.2f}")
+    return hit_at_one
+
+
+def compute_hl_ap(qid2preds, qid2gt_scores_binary, num_workers=8, chunksize=50):
+    qid2pred_scores = {k: v["pred_saliency_scores"] for k, v in qid2preds.items()}
+    ap_scores = np.zeros((len(qid2preds), 3))   # (#preds, 3)
+    qids = list(qid2preds.keys())
+    input_tuples = []
+    for idx, qid in enumerate(qids):
+        for w_idx in range(3):  # annotation score idx
+            y_true = qid2gt_scores_binary[qid][:, w_idx]
+            y_predict = np.array(qid2pred_scores[qid])
+            input_tuples.append((idx, w_idx, y_true, y_predict))
+
+    if num_workers > 1:
+        with mp.Pool(num_workers) as pool:
+            for idx, w_idx, score in pool.imap_unordered(
+                    compute_ap_from_tuple, input_tuples, chunksize=chunksize):
+                ap_scores[idx, w_idx] = score
+    else:
+        for input_tuple in input_tuples:
+            idx, w_idx, score = compute_ap_from_tuple(input_tuple)
+            ap_scores[idx, w_idx] = score
+
+    # it's the same if we first average across different annotations, then average across queries
+    # since all queries have the same #annotations.
+    mean_ap = float(f"{100 * np.mean(ap_scores):.2f}")
+    return mean_ap
+
+
+def compute_ap_from_tuple(input_tuple):
+    idx, w_idx, y_true, y_predict = input_tuple
+    if len(y_true) < len(y_predict):
+        # print(f"len(y_true) < len(y_predict) {len(y_true), len(y_predict)}")
+        y_predict = y_predict[:len(y_true)]
+    elif len(y_true) > len(y_predict):
+        # print(f"len(y_true) > len(y_predict) {len(y_true), len(y_predict)}")
+        _y_predict = np.zeros(len(y_true))
+        _y_predict[:len(y_predict)] = y_predict
+        y_predict = _y_predict
+
+    score = get_ap(y_true, y_predict)
+    return idx, w_idx, score
+
+
+def mk_gt_scores(gt_data, clip_length=2):
+    """gt_data, dict, """
+    num_clips = int(gt_data["duration"] / clip_length)
+    saliency_scores_full_video = np.zeros((num_clips, 3))
+    relevant_clip_ids = np.array(gt_data["relevant_clip_ids"])  # (#relevant_clip_ids, )
+    saliency_scores_relevant_clips = np.array(gt_data["saliency_scores"])  # (#relevant_clip_ids, 3)
+    saliency_scores_full_video[relevant_clip_ids] = saliency_scores_relevant_clips
+    return saliency_scores_full_video  # (#clips_in_video, 3)  the scores are in range [0, 4]
+
+
+def eval_highlight(submission, ground_truth, verbose=True):
+    """
+    Args:
+        submission:
+        ground_truth:
+        verbose:
+    """
+    qid2preds = {d["qid"]: d for d in submission}
+    qid2gt_scores_full_range = {d["qid"]: mk_gt_scores(d) for d in ground_truth}  # scores in range [0, 4]
+    # gt_saliency_score_min: int, in [0, 1, 2, 3, 4]. The minimum score for a positive clip.
+    gt_saliency_score_min_list = [2, 3, 4]
+    saliency_score_names = ["Fair", "Good", "VeryGood"]
+    highlight_det_metrics = {}
+    for gt_saliency_score_min, score_name in zip(gt_saliency_score_min_list, saliency_score_names):
+        start_time = time.time()
+        qid2gt_scores_binary = {
+            k: (v >= gt_saliency_score_min).astype(float)
+            for k, v in qid2gt_scores_full_range.items()}  # scores in [0, 1]
+        hit_at_one = compute_hl_hit1(qid2preds, qid2gt_scores_binary)
+        mean_ap = compute_hl_ap(qid2preds, qid2gt_scores_binary)
+        highlight_det_metrics[f"HL-min-{score_name}"] = {"HL-mAP": mean_ap, "HL-Hit1": hit_at_one}
+        if verbose:
+            print(f"Calculating highlight scores with min score {gt_saliency_score_min} ({score_name})")
+            print(f"Time cost {time.time() - start_time:.2f} seconds")
+    return highlight_det_metrics
+
+
+def eval_submission(submission, ground_truth, verbose=True, match_number=False, hl=False):
+    """
+    Args:
+        submission: list(dict), each dict is {
+            qid: str,
+            query: str,
+            vid: str,
+            pred_relevant_windows: list([st, ed]),
+            pred_saliency_scores: list(float), len == #clips in video.
+                i.e., each clip in the video will have a saliency score.
+        }
+        ground_truth: list(dict), each dict is     {
+          "qid": 7803,
+          "query": "Man in gray top walks from outside to inside.",
+          "duration": 150,
+          "vid": "RoripwjYFp8_360.0_510.0",
+          "relevant_clip_ids": [13, 14, 15, 16, 17]
+          "saliency_scores": [[4, 4, 2], [3, 4, 2], [2, 2, 3], [2, 2, 2], [0, 1, 3]]
+               each sublist corresponds to one clip in relevant_clip_ids.
+               The 3 elements in the sublist are scores from 3 different workers. The
+               scores are in [0, 1, 2, 3, 4], meaning [Very Bad, ..., Good, Very Good]
+        }
+        verbose:
+        match_number:
+
+    Returns:
+
+    """
+    pred_qids = set([e["qid"] for e in submission])
+    gt_qids = set([e["qid"] for e in ground_truth])
+    # import pdb; pdb.set_trace()
+    if match_number:
+        assert pred_qids == gt_qids, \
+            f"qids in ground_truth and submission must match. " \
+            f"use `match_number=False` if you wish to disable this check"
+    else:  # only leave the items that exists in both submission and ground_truth
+        shared_qids = pred_qids.intersection(gt_qids)
+        submission = [e for e in submission if e["qid"] in shared_qids]
+        ground_truth = [e for e in ground_truth if e["qid"] in shared_qids]
+
+    eval_metrics = {}
+    eval_metrics_brief = OrderedDict()
+    if "pred_relevant_windows" in submission[0]:
+        moment_ret_scores = eval_moment_retrieval(submission, ground_truth, verbose=verbose)
+        eval_metrics.update(moment_ret_scores)
+        moment_ret_scores_brief = {
+            "MR-full-mAP": moment_ret_scores["full"]["MR-mAP"]["average"],
+            "[email protected]": moment_ret_scores["full"]["MR-mAP"]["0.5"],
+            "[email protected]": moment_ret_scores["full"]["MR-mAP"]["0.75"],
+            "MR-short-mAP": moment_ret_scores["short"]["MR-mAP"]["average"],
+            "MR-middle-mAP": moment_ret_scores["middle"]["MR-mAP"]["average"],
+            "MR-long-mAP": moment_ret_scores["long"]["MR-mAP"]["average"],
+            "MR-full-mIoU": moment_ret_scores["full"]["MR-mIoU"],
+            "[email protected]": moment_ret_scores["full"]["MR-R1"]["0.3"],
+            "[email protected]": moment_ret_scores["full"]["MR-R1"]["0.5"],
+            "[email protected]": moment_ret_scores["full"]["MR-R1"]["0.7"],
+        }
+        eval_metrics_brief.update(
+            sorted([(k, v) for k, v in moment_ret_scores_brief.items()], key=lambda x: x[0]))
+
+    if "pred_saliency_scores" in submission[0] and hl:
+        highlight_det_scores = eval_highlight(
+            submission, ground_truth, verbose=verbose)
+        eval_metrics.update(highlight_det_scores)
+        highlight_det_scores_brief = dict([
+            (f"{k}-{sub_k.split('-')[1]}", v[sub_k])
+            for k, v in highlight_det_scores.items() for sub_k in v])
+        eval_metrics_brief.update(highlight_det_scores_brief)
+
+    # sort by keys
+    final_eval_metrics = OrderedDict()
+    final_eval_metrics["brief"] = eval_metrics_brief
+    final_eval_metrics.update(sorted([(k, v) for k, v in eval_metrics.items()], key=lambda x: x[0]))
+    return final_eval_metrics
+
+
+def eval_main():
+    import argparse
+    parser = argparse.ArgumentParser(description="Moments and Highlights Evaluation Script")
+    parser.add_argument("--submission_path", type=str, help="path to generated prediction file")
+    parser.add_argument("--gt_path", type=str, help="path to GT file")
+    parser.add_argument("--save_path", type=str, help="path to save the results")
+    parser.add_argument("--not_verbose", action="store_true")
+    args = parser.parse_args()
+
+    verbose = not args.not_verbose
+    submission = load_jsonl(args.submission_path)
+    gt = load_jsonl(args.gt_path)
+    results = eval_submission(submission, gt, verbose=verbose)
+    if verbose:
+        print(json.dumps(results, indent=4))
+
+    with open(args.save_path, "w") as f:
+        f.write(json.dumps(results, indent=4))
+
+
+if __name__ == '__main__':
+    eval_main()

third_party/cgdetr/standalone_eval/eval_sample.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+# Usage: bash standalone_eval/eval_sample.sh
+submission_path=standalone_eval/sample_val_preds.jsonl
+gt_path=data/highlight_val_release.jsonl
+save_path=standalone_eval/sample_val_preds_metrics.json
+
+PYTHONPATH=$PYTHONPATH:. python standalone_eval/eval.py \
+--submission_path ${submission_path} \
+--gt_path ${gt_path} \
+--save_path ${save_path}

third_party/cgdetr/standalone_eval/utils.py

@@ -0,0 +1,209 @@
+"""
+Copied from MMAction2
+https://github.com/open-mmlab/mmaction2/blob/master/mmaction/core/evaluation/eval_detection.py
+"""
+import json
+import numpy as np
+from sklearn.metrics import precision_recall_curve
+
+
+def load_jsonl(filename):
+    with open(filename, "r") as f:
+        return [json.loads(l.strip("\n")) for l in f.readlines()]
+
+
+def compute_temporal_iou_batch_paired(pred_windows, gt_windows):
+    """ compute intersection-over-union along temporal axis for each pair of windows in pred_windows and gt_windows.
+    Args:
+        pred_windows: np.ndarray, (N, 2), [st (float), ed (float)] * N
+        gt_windows: np.ndarray, (N, 2), [st (float), ed (float)] * N
+    Returns:
+        iou (float): np.ndarray, (N, )
+
+    References:
+        for np.divide with zeros, see https://stackoverflow.com/a/37977222
+    """
+    intersection = np.maximum(
+        0, np.minimum(pred_windows[:, 1], gt_windows[:, 1]) - np.maximum(pred_windows[:, 0], gt_windows[:, 0])
+    )
+    union = np.maximum(pred_windows[:, 1], gt_windows[:, 1]) \
+            - np.minimum(pred_windows[:, 0], gt_windows[:, 0])  # not the correct union though
+    return np.divide(intersection, union, out=np.zeros_like(intersection), where=union != 0)
+
+
+def compute_temporal_iou_batch_cross(spans1, spans2):
+    """
+    Args:
+        spans1: (N, 2) np.ndarray, each row defines a span [st, ed]
+        spans2: (M, 2) np.ndarray, ...
+
+    Returns:
+        iou: (N, M) np.ndarray
+        union: (N, M) np.ndarray
+    >>> spans1 = np.array([[0, 0.2, 0.9], [0.5, 1.0, 0.2]])
+    >>> spans2 = np.array([[0, 0.3], [0., 1.0]])
+    >>> compute_temporal_iou_batch_cross(spans1, spans2)
+    (tensor([[0.6667, 0.2000],
+         [0.0000, 0.5000]]),
+     tensor([[0.3000, 1.0000],
+             [0.8000, 1.0000]]))
+    """
+    areas1 = spans1[:, 1] - spans1[:, 0]  # (N, )
+    areas2 = spans2[:, 1] - spans2[:, 0]  # (M, )
+
+    left = np.maximum(spans1[:, None, 0], spans2[None, :, 0])  # (N, M)
+    right = np.minimum(spans1[:, None, 1], spans2[None, :, 1])  # (N, M)
+
+    inter = np.clip(right - left, 0, None)  # (N, M)
+    union = areas1[:, None] + areas2[None, :] - inter  # (N, M)
+
+    iou = inter / union
+    return iou, union
+
+
+def interpolated_precision_recall(precision, recall):
+    """Interpolated AP - VOCdevkit from VOC 2011.
+
+    Args:
+        precision (np.ndarray): The precision of different thresholds.
+        recall (np.ndarray): The recall of different thresholds.
+
+    Returns：
+        float: Average precision score.
+    """
+    mprecision = np.hstack([[0], precision, [0]])
+    mrecall = np.hstack([[0], recall, [1]])
+    for i in range(len(mprecision) - 1)[::-1]:
+        mprecision[i] = max(mprecision[i], mprecision[i + 1])
+    idx = np.where(mrecall[1::] != mrecall[0:-1])[0] + 1
+    ap = np.sum((mrecall[idx] - mrecall[idx - 1]) * mprecision[idx])
+    return ap
+
+
+def compute_average_precision_detection(ground_truth,
+                                        prediction,
+                                        tiou_thresholds=np.linspace(
+                                            0.5, 0.95, 10)):
+    """Compute average precision (detection task) between ground truth and
+    predictions data frames. If multiple predictions occurs for the same
+    predicted segment, only the one with highest score is matches as true
+    positive. This code is greatly inspired by Pascal VOC devkit.
+
+    Args:
+        ground_truth (list[dict]): List containing the ground truth instances
+            (dictionaries). Required keys are 'video-id', 't-start' and
+            't-end'.
+        prediction (list[dict]): List containing the prediction instances
+            (dictionaries). Required keys are: 'video-id', 't-start', 't-end'
+            and 'score'.
+        tiou_thresholds (np.ndarray): A 1darray indicates the temporal
+            intersection over union threshold, which is optional.
+            Default: ``np.linspace(0.5, 0.95, 10)``.
+
+    Returns:
+        Float: ap, Average precision score.
+    """
+    num_thresholds = len(tiou_thresholds)
+    num_gts = len(ground_truth)
+    num_preds = len(prediction)
+    ap = np.zeros(num_thresholds)
+    if len(prediction) == 0:
+        return ap
+
+    num_positive = float(num_gts)
+    lock_gt = np.ones((num_thresholds, num_gts)) * -1
+    # Sort predictions by decreasing score order.
+    prediction.sort(key=lambda x: -x['score'])
+    # Initialize true positive and false positive vectors.
+    tp = np.zeros((num_thresholds, num_preds))
+    fp = np.zeros((num_thresholds, num_preds))
+
+    # Adaptation to query faster
+    ground_truth_by_videoid = {}
+    for i, item in enumerate(ground_truth):
+        item['index'] = i
+        ground_truth_by_videoid.setdefault(item['video-id'], []).append(item)
+
+    # Assigning true positive to truly grount truth instances.
+    for idx, pred in enumerate(prediction):
+        if pred['video-id'] in ground_truth_by_videoid:
+            gts = ground_truth_by_videoid[pred['video-id']]
+        else:
+            fp[:, idx] = 1
+            continue
+
+        _pred = np.array([[pred['t-start'], pred['t-end']], ])
+        _gt = np.array([[gt['t-start'], gt['t-end']] for gt in gts])
+        tiou_arr = compute_temporal_iou_batch_cross(_pred, _gt)[0]
+
+        tiou_arr = tiou_arr.reshape(-1)
+        # We would like to retrieve the predictions with highest tiou score.
+        tiou_sorted_idx = tiou_arr.argsort()[::-1]
+        for t_idx, tiou_threshold in enumerate(tiou_thresholds):
+            for j_idx in tiou_sorted_idx:
+                if tiou_arr[j_idx] < tiou_threshold:
+                    fp[t_idx, idx] = 1
+                    break
+                if lock_gt[t_idx, gts[j_idx]['index']] >= 0:
+                    continue
+                # Assign as true positive after the filters above.
+                tp[t_idx, idx] = 1
+                lock_gt[t_idx, gts[j_idx]['index']] = idx
+                break
+
+            if fp[t_idx, idx] == 0 and tp[t_idx, idx] == 0:
+                fp[t_idx, idx] = 1
+
+    tp_cumsum = np.cumsum(tp, axis=1).astype(float)
+    fp_cumsum = np.cumsum(fp, axis=1).astype(float)
+    recall_cumsum = tp_cumsum / num_positive
+
+    precision_cumsum = tp_cumsum / (tp_cumsum + fp_cumsum)
+
+    for t_idx in range(len(tiou_thresholds)):
+        ap[t_idx] = interpolated_precision_recall(precision_cumsum[t_idx, :],
+                                                  recall_cumsum[t_idx, :])
+    return ap
+
+
+def get_ap(y_true, y_predict, interpolate=True, point_11=False):
+    """
+    Average precision in different formats: (non-) interpolated and/or 11-point approximated
+    point_11=True and interpolate=True corresponds to the 11-point interpolated AP used in
+    the PASCAL VOC challenge up to the 2008 edition and has been verfied against the vlfeat implementation
+    The exact average precision (interpolate=False, point_11=False) corresponds to the one of vl_feat
+
+    :param y_true: list/ numpy vector of true labels in {0,1} for each element
+    :param y_predict: predicted score for each element
+    :param interpolate: Use interpolation?
+    :param point_11: Use 11-point approximation to average precision?
+    :return: average precision
+
+    ref: https://github.com/gyglim/video2gif_dataset/blob/master/v2g_evaluation/__init__.py
+
+    """
+    # Check inputs
+    assert len(y_true) == len(y_predict), "Prediction and ground truth need to be of the same length"
+    if len(set(y_true)) == 1:
+        if y_true[0] == 0:
+            return 0  # True labels are all zeros
+            # raise ValueError('True labels cannot all be zero')
+        else:
+            return 1
+    else:
+        assert sorted(set(y_true)) == [0, 1], "Ground truth can only contain elements {0,1}"
+
+    # Compute precision and recall
+    precision, recall, _ = precision_recall_curve(y_true, y_predict)
+    recall = recall.astype(np.float32)
+
+    if interpolate:  # Compute the interpolated precision
+        for i in range(1, len(precision)):
+            precision[i] = max(precision[i - 1], precision[i])
+
+    if point_11:  # Compute the 11-point approximated AP
+        precision_11 = [precision[np.where(recall >= t)[0][-1]] for t in np.arange(0, 1.01, 0.1)]
+        return np.mean(precision_11)
+    else:  # Compute the AP using precision at every additionally recalled sample
+        indices = np.where(np.diff(recall))
+        return np.mean(precision[indices])

third_party/cgdetr/utils/basic_utils.py

@@ -0,0 +1,221 @@
+import os
+import json
+import zipfile
+import numpy as np
+import pickle
+from collections import OrderedDict, Counter
+import pandas as pd
+
+
+def load_pickle(filename):
+    with open(filename, "rb") as f:
+        return pickle.load(f)
+
+
+def save_pickle(data, filename):
+    with open(filename, "wb") as f:
+        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)
+
+
+def load_json(filename):
+    with open(filename, "r") as f:
+        return json.load(f)
+
+
+def save_json(data, filename, save_pretty=False, sort_keys=False):
+    with open(filename, "w") as f:
+        if save_pretty:
+            f.write(json.dumps(data, indent=4, sort_keys=sort_keys))
+        else:
+            json.dump(data, f)
+
+
+def load_jsonl(filename):
+    with open(filename, "r") as f:
+        return [json.loads(l.strip("\n")) for l in f.readlines()]
+
+
+def save_jsonl(data, filename):
+    """data is a list"""
+    with open(filename, "w") as f:
+        f.write("\n".join([json.dumps(e) for e in data]))
+
+
+def save_lines(list_of_str, filepath):
+    with open(filepath, "w") as f:
+        f.write("\n".join(list_of_str))
+
+
+def read_lines(filepath):
+    with open(filepath, "r") as f:
+        return [e.strip("\n") for e in f.readlines()]
+
+
+def mkdirp(p):
+    if not os.path.exists(p):
+        os.makedirs(p)
+
+
+def flat_list_of_lists(l):
+    """flatten a list of lists [[1,2], [3,4]] to [1,2,3,4]"""
+    return [item for sublist in l for item in sublist]
+
+
+def convert_to_seconds(hms_time):
+    """ convert '00:01:12' to 72 seconds.
+    :hms_time (str): time in comma separated string, e.g. '00:01:12'
+    :return (int): time in seconds, e.g. 72
+    """
+    times = [float(t) for t in hms_time.split(":")]
+    return times[0] * 3600 + times[1] * 60 + times[2]
+
+
+def get_video_name_from_url(url):
+    return url.split("/")[-1][:-4]
+
+
+def merge_dicts(list_dicts):
+    merged_dict = list_dicts[0].copy()
+    for i in range(1, len(list_dicts)):
+        merged_dict.update(list_dicts[i])
+    return merged_dict
+
+
+def l2_normalize_np_array(np_array, eps=1e-5):
+    """np_array: np.ndarray, (*, D), where the last dim will be normalized"""
+    return np_array / (np.linalg.norm(np_array, axis=-1, keepdims=True) + eps)
+
+
+def make_zipfile(src_dir, save_path, enclosing_dir="", exclude_dirs=None, exclude_extensions=None,
+                 exclude_dirs_substring=None):
+    """make a zip file of root_dir, save it to save_path.
+    exclude_paths will be excluded if it is a subdir of root_dir.
+    An enclosing_dir is added is specified.
+    """
+    abs_src = os.path.abspath(src_dir)
+    with zipfile.ZipFile(save_path, "w") as zf:
+        for dirname, subdirs, files in os.walk(src_dir):
+            if exclude_dirs is not None:
+                for e_p in exclude_dirs:
+                    if e_p in subdirs:
+                        subdirs.remove(e_p)
+            if exclude_dirs_substring is not None:
+                to_rm = []
+                for d in subdirs:
+                    if exclude_dirs_substring in d:
+                        to_rm.append(d)
+                for e in to_rm:
+                    subdirs.remove(e)
+            arcname = os.path.join(enclosing_dir, dirname[len(abs_src) + 1:])
+            zf.write(dirname, arcname)
+            for filename in files:
+                if exclude_extensions is not None:
+                    if os.path.splitext(filename)[1] in exclude_extensions:
+                        continue  # do not zip it
+                absname = os.path.join(dirname, filename)
+                arcname = os.path.join(enclosing_dir, absname[len(abs_src) + 1:])
+                zf.write(absname, arcname)
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current/max/min value"""
+    def __init__(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+        self.max = -1e10
+        self.min = 1e10
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+        self.max = -1e10
+        self.min = 1e10
+
+    def update(self, val, n=1):
+        self.max = max(val, self.max)
+        self.min = min(val, self.min)
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def dissect_by_lengths(np_array, lengths, dim=0, assert_equal=True):
+    """Dissect an array (N, D) into a list a sub-array,
+    np_array.shape[0] == sum(lengths), Output is a list of nd arrays, singlton dimention is kept"""
+    if assert_equal:
+        assert len(np_array) == sum(lengths)
+    length_indices = [0, ]
+    for i in range(len(lengths)):
+        length_indices.append(length_indices[i] + lengths[i])
+    if dim == 0:
+        array_list = [np_array[length_indices[i]:length_indices[i+1]] for i in range(len(lengths))]
+    elif dim == 1:
+        array_list = [np_array[:, length_indices[i]:length_indices[i + 1]] for i in range(len(lengths))]
+    elif dim == 2:
+        array_list = [np_array[:, :, length_indices[i]:length_indices[i + 1]] for i in range(len(lengths))]
+    else:
+        raise NotImplementedError
+    return array_list
+
+
+def get_ratio_from_counter(counter_obj, threshold=200):
+    keys = counter_obj.keys()
+    values = counter_obj.values()
+    filtered_values = [counter_obj[k] for k in keys if k > threshold]
+    return float(sum(filtered_values)) / sum(values)
+
+
+def get_counter_dist(counter_object, sort_type="none"):
+    _sum = sum(counter_object.values())
+    dist = {k: float(f"{100 * v / _sum:.2f}") for k, v in counter_object.items()}
+    if sort_type == "value":
+        dist = OrderedDict(sorted(dist.items(), reverse=True))
+    return dist
+
+
+def get_show_name(vid_name):
+    """
+    get tvshow name from vid_name
+    :param vid_name: video clip name
+    :return: tvshow name
+    """
+    show_list = ["friends", "met", "castle", "house", "grey"]
+    vid_name_prefix = vid_name.split("_")[0]
+    show_name = vid_name_prefix if vid_name_prefix in show_list else "bbt"
+    return show_name
+
+
+def get_abspaths_by_ext(dir_path, ext=(".jpg",)):
+    """Get absolute paths to files in dir_path with extensions specified by ext.
+    Note this function does work recursively.
+    """
+    if isinstance(ext, list):
+        ext = tuple(ext)
+    if isinstance(ext, str):
+        ext = tuple([ext, ])
+    filepaths = [os.path.join(root, name)
+                 for root, dirs, files in os.walk(dir_path)
+                 for name in files
+                 if name.endswith(tuple(ext))]
+    return filepaths
+
+
+def get_basename_no_ext(path):
+    """ '/data/movienet/240p_keyframe_feats/tt7672188.npz' --> 'tt7672188' """
+    return os.path.splitext(os.path.split(path)[1])[0]
+
+
+def dict_to_markdown(d, max_str_len=120):
+    # convert list into its str representation
+    d = {k: v.__repr__() if isinstance(v, list) else v for k, v in d.items()}
+    # truncate string that is longer than max_str_len
+    if max_str_len is not None:
+        d = {k: v[-max_str_len:] if isinstance(v, str) else v for k, v in d.items()}
+    return pd.DataFrame(d, index=[0]).transpose().to_markdown()
+