Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+videos/3249402410.mp4 filter=lfs diff=lfs merge=lfs -text
+videos/4882821564.mp4 filter=lfs diff=lfs merge=lfs -text
+videos/6233408665.mp4 filter=lfs diff=lfs merge=lfs -text

.github/workflows/update_space.yml

@@ -0,0 +1,28 @@
+name: Run Python script
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Install Gradio
+      run: python -m pip install gradio
+
+    - name: Log in to Hugging Face
+      run: python -c 'import huggingface_hub; huggingface_hub.login(token="${{ secrets.hf_token }}")'
+
+    - name: Deploy to Spaces
+      run: gradio deploy

Infer.py

@@ -0,0 +1,149 @@
+
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+import math
+from tqdm import tqdm
+import argparse
+from collections import OrderedDict
+import json
+
+from collections import defaultdict
+from model.deberta_moe import DebertaV2ForMaskedLM
+from transformers import DebertaV2Tokenizer
+
+import clip
+import ffmpeg
+from VideoLoader import VideoLoader
+
+def get_mask(lengths, max_length):
+    """ Computes a batch of padding masks given batched lengths """
+    mask = 1 * (
+        torch.arange(max_length).unsqueeze(1) < lengths
+    ).transpose(0, 1)
+    return mask
+
+class Infer:
+    def __init__(self, device):
+        pretrained_ckpt = torch.load("ckpts/model.pth")
+        args = pretrained_ckpt['args']
+        args.n_ans = 2
+        args.max_tokens = 256
+        self.args = args
+        self.clip_model = clip.load("ViT-L/14", device = device)[0]
+        self.tokenizer = DebertaV2Tokenizer.from_pretrained(
+            "ckpts/deberta-v2-xlarge", local_files_only=True
+        )
+        
+        self.model = DebertaV2ForMaskedLM.from_pretrained(
+                        features_dim=args.features_dim if args.use_video else 0,
+                        max_feats=args.max_feats,
+                        freeze_lm=args.freeze_lm,
+                        freeze_mlm=args.freeze_mlm,
+                        ft_ln=args.ft_ln,
+                        ds_factor_attn=args.ds_factor_attn,
+                        ds_factor_ff=args.ds_factor_ff,
+                        dropout=args.dropout,
+                        n_ans=args.n_ans,
+                        freeze_last=args.freeze_last,
+                        pretrained_model_name_or_path="ckpts/deberta-v2-xlarge",
+                        local_files_only=False,
+                        add_video_feat=args.add_video_feat,
+                        freeze_ad=args.freeze_ad,
+                    )
+        new_state_dict = OrderedDict()
+        for k, v in pretrained_ckpt['model'].items():
+            new_state_dict[k.replace("module.","")] = v
+        self.model.load_state_dict(pretrained_ckpt, strict=False)
+        self.model.eval()
+        self.model.to(device)
+        self.device = device
+
+        self.video_loader = VideoLoader()
+        self.set_answer()
+
+    def _get_clip_feature(self, video):
+        feat = self.clip_model.encode_image(video.to(self.device))
+        #feat = F.normalize(feat, dim=1)
+        return feat
+
+    def set_answer(self):
+        tok_yes = torch.tensor(
+                    self.tokenizer(
+                        "Yes",
+                        add_special_tokens=False,
+                        max_length=1,
+                        truncation=True,
+                        padding="max_length",
+                    )["input_ids"],
+                    dtype=torch.long,
+                )
+        tok_no = torch.tensor(
+            self.tokenizer(
+                "No",
+                add_special_tokens=False,
+                max_length=1,
+                truncation=True,
+                padding="max_length",
+            )["input_ids"],
+            dtype=torch.long,
+        )     
+
+        a2tok = torch.stack([tok_yes, tok_no])
+        self.model.set_answer_embeddings(
+            a2tok.to(self.model.device), freeze_last=self.args.freeze_last
+        )
+
+    def generate(self, text, video_path, candidates = None):
+        video, video_len = self.video_loader(video_path)
+        video = self._get_clip_feature(video).unsqueeze(0).float()
+        video_mask = get_mask(video_len, 10)
+        video_mask = torch.cat([torch.ones((1,1)),video_mask], dim=1)
+        logits_list = []
+        
+        question = text.capitalize().strip()
+        if question[-1] != "?":
+            question = str(question) + "?"
+
+        for aid in range(len(candidates)):
+            prompt = (
+                f" Question: {question} Is it '{candidates[aid]}'? {self.tokenizer.mask_token}. Subtitles: "
+            )
+            prompt = prompt.strip()
+            encoded = self.tokenizer(
+                prompt,
+                add_special_tokens=True,
+                max_length=self.args.max_tokens,
+                padding="longest",
+                truncation=True,
+                return_tensors="pt",
+            )
+            # forward
+
+            output = self.model(
+                video=video.to(self.device),
+                video_mask=video_mask.to(self.device),
+                input_ids=encoded["input_ids"].to(self.device),
+                attention_mask=encoded["attention_mask"].to(self.device),
+            )
+            # += output['loads'].detach().cpu()
+            logits = output["logits"]
+            # get logits for the mask token
+            delay = 11
+            logits = logits[:, delay : encoded["input_ids"].size(1) + delay][
+                encoded["input_ids"] == self.tokenizer.mask_token_id
+            ]
+            logits_list.append(logits.softmax(-1)[:, 0])
+        
+        logits = torch.stack(logits_list, 1)
+        if logits.shape[1] == 1:
+            preds = logits.round().long().squeeze(1)
+        else:
+            preds = logits.max(1).indices
+
+        return candidates[preds]
+    

README.md

@@ -1,13 +1,6 @@
 ---
-title: T MoENet
-emoji: 😻
-colorFrom: yellow
-colorTo: green
-sdk: gradio
-sdk_version: 4.38.1
+title: T-MoENet
 app_file: app.py
-pinned: false
-license: apache-2.0
+sdk: gradio
+sdk_version: 3.46.0
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

VideoLoader.py

@@ -0,0 +1,133 @@
+
+import torch as th
+import os
+import numpy as np
+import ffmpeg
+
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = th.FloatTensor(mean).view(1, 3, 1, 1)
+        self.std = th.FloatTensor(std).view(1, 3, 1, 1)
+
+    def __call__(self, tensor):
+        tensor = (tensor - self.mean) / (self.std + 1e-8)
+        return tensor
+
+
+class Preprocessing(object):
+    def __init__(self):
+        self.norm = Normalize(
+            mean=[0.48145466, 0.4578275, 0.40821073],
+            std=[0.26862954, 0.26130258, 0.27577711],
+        )
+
+    def __call__(self, tensor):
+        tensor = tensor / 255.0
+        tensor = self.norm(tensor)
+        return tensor
+
+
+class VideoLoader:
+    """Pytorch video loader."""
+
+    def __init__(
+        self,
+        framerate=1,
+        size=224,
+        centercrop=True,
+    ):
+        self.centercrop = centercrop
+        self.size = size
+        self.framerate = framerate
+        self.preprocess = Preprocessing()
+        self.max_feats = 10
+        self.features_dim = 768
+
+    def _get_video_dim(self, video_path):
+        probe = ffmpeg.probe(video_path)
+        video_stream = next(
+            (stream for stream in probe["streams"] if stream["codec_type"] == "video"),
+            None,
+        )
+        width = int(video_stream["width"])
+        height = int(video_stream["height"])
+        num, denum = video_stream["avg_frame_rate"].split("/")
+        frame_rate = int(num) / int(denum)
+        return height, width, frame_rate
+
+    def _get_output_dim(self, h, w):
+        if isinstance(self.size, tuple) and len(self.size) == 2:
+            return self.size
+        elif h >= w:
+            return int(h * self.size / w), self.size
+        else:
+            return self.size, int(w * self.size / h)
+
+    def _getvideo(self, video_path):
+        
+        if os.path.isfile(video_path):
+            print("Decoding video: {}".format(video_path))
+            try:
+                h, w, fr = self._get_video_dim(video_path)
+            except:
+                print("ffprobe failed at: {}".format(video_path))
+                return {
+                    "video": th.zeros(1),
+                    "input": video_path
+                }
+            if fr < 1:
+                print("Corrupted Frame Rate: {}".format(video_path))
+                return {
+                    "video": th.zeros(1),
+                    "input": video_path
+                }
+            height, width = self._get_output_dim(h, w)
+
+            try:
+                cmd = (
+                    ffmpeg.input(video_path)
+                    .filter("fps", fps=self.framerate)
+                    .filter("scale", width, height)
+                )
+                if self.centercrop:
+                    x = int((width - self.size) / 2.0)
+                    y = int((height - self.size) / 2.0)
+                    cmd = cmd.crop(x, y, self.size, self.size)
+                out, _ = cmd.output("pipe:", format="rawvideo", pix_fmt="rgb24").run(
+                    capture_stdout=True, quiet=True
+                )
+            except:
+                print("ffmpeg error at: {}".format(video_path))
+                return {
+                    "video": th.zeros(1),
+                    "input": video_path,
+                }
+            if self.centercrop and isinstance(self.size, int):
+                height, width = self.size, self.size
+            video = np.frombuffer(out, np.uint8).reshape([-1, height, width, 3])
+            video = th.from_numpy(video.astype("float32"))
+            video = video.permute(0, 3, 1, 2) # t,c,h,w
+        else:
+            video = th.zeros(1)
+
+        return {"video": video, "input": video_path}
+
+    def __call__(self, video_path):
+
+        video = self._getvideo(video_path)['video']
+
+        if len(video) > self.max_feats:
+            sampled = []
+            for j in range(self.max_feats):
+                sampled.append(video[(j * len(video)) // self.max_feats])
+            video = th.stack(sampled)
+            video_len = self.max_feats
+        elif len(video) < self.max_feats:
+            video_len = len(video)
+            video = th.cat(
+                [video, th.zeros(self.max_feats - video_len, self.features_dim)], 0
+            )
+        video = self.preprocess(video)
+        return video, video_len
+    

app.py

@@ -0,0 +1,157 @@
+import shutil
+import gradio as gr
+import torch
+from fastapi import FastAPI
+import os
+import tempfile
+from Infer import Infer
+
+title_markdown = ("""
+<div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+  <div>
+    <h1 >Temporal-guided Mixture-of-Experts for Zero-Shot Video Question Answering</h1>
+    <h5 style="margin: 0;">Under review.</h5>
+  </div>
+</div>
+                  
+<div align="center">
+    <div style="display:flex; gap: 0.25rem;" align="center">
+        <a href='https://github.com/qyx1121/T-MoENet'><img src='https://img.shields.io/badge/Github-Code-blue'></a>
+    </div>
+</div>
+""")
+
+block_css = """
+#buttons button {
+    min-width: min(120px,100%);
+}
+"""
+
+def save_video_to_local(video_path):
+    filename = os.path.join('temp', next(tempfile._get_candidate_names()) + '.mp4')
+    shutil.copyfile(video_path, filename)
+    return filename
+
+
+def generate(video, textbox_in, first_run, state, state_):
+    flag = 1
+    if not textbox_in:
+        if len(state_.messages) > 0:
+            textbox_in = state_.messages[-1][1]
+            state_.messages.pop(-1)
+            flag = 0
+        else:
+            return "Please enter instruction"
+    video = video if video else "none"
+    # assert not (os.path.exists(image1) and os.path.exists(video))
+
+    first_run = False if len(state.messages) > 0 else True
+
+    text_en_in = textbox_in.replace("picture", "image")
+
+    # images_tensor = [[], []]
+    image_processor = handler.image_processor
+    if os.path.exists(image1) and not os.path.exists(video):
+        tensor = image_processor.preprocess(image1, return_tensors='pt')['pixel_values'][0]
+        # print(tensor.shape)
+        tensor = tensor.to(handler.model.device, dtype=dtype)
+        images_tensor[0] = images_tensor[0] + [tensor]
+        images_tensor[1] = images_tensor[1] + ['image']
+        print(torch.cuda.memory_allocated())
+        print(torch.cuda.max_memory_allocated())
+    video_processor = handler.video_processor
+    if not os.path.exists(image1) and os.path.exists(video):
+        tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]
+        # print(tensor.shape)
+        tensor = tensor.to(handler.model.device, dtype=dtype)
+        images_tensor[0] = images_tensor[0] + [tensor]
+        images_tensor[1] = images_tensor[1] + ['video']
+        print(torch.cuda.memory_allocated())
+        print(torch.cuda.max_memory_allocated())
+    if os.path.exists(image1) and os.path.exists(video):
+        tensor = video_processor(video, return_tensors='pt')['pixel_values'][0]
+        # print(tensor.shape)
+        tensor = tensor.to(handler.model.device, dtype=dtype)
+        images_tensor[0] = images_tensor[0] + [tensor]
+        images_tensor[1] = images_tensor[1] + ['video']
+        
+        
+        tensor = image_processor.preprocess(image1, return_tensors='pt')['pixel_values'][0]
+        # print(tensor.shape)
+        tensor = tensor.to(handler.model.device, dtype=dtype)
+        images_tensor[0] = images_tensor[0] + [tensor]
+        images_tensor[1] = images_tensor[1] + ['image']
+        print(torch.cuda.memory_allocated())
+        print(torch.cuda.max_memory_allocated())
+        
+
+    text_en_out, state_ = handler.generate(images_tensor, text_en_in, first_run=first_run, state=state_)
+    state_.messages[-1] = (state_.roles[1], text_en_out)
+
+    text_en_out = text_en_out.split('#')[0]
+    textbox_out = text_en_out
+
+    show_images = ""
+    if flag:
+        state.append_message(state.roles[0], textbox_in + "\n" + show_images)
+    state.append_message(state.roles[1], textbox_out)
+    torch.cuda.empty_cache()
+    return (state, state_, state.to_gradio_chatbot(), False, gr.update(value=None, interactive=True), images_tensor, gr.update(value=image1 if os.path.exists(image1) else None, interactive=True), gr.update(value=video if os.path.exists(video) else None, interactive=True))
+
+
+device = "cuda"
+handler = Infer(device)
+# handler.model.to(dtype=dtype)
+if not os.path.exists("temp"):
+    os.makedirs("temp")
+
+print(torch.cuda.memory_allocated())
+print(torch.cuda.max_memory_allocated())
+
+textbox = gr.Textbox(
+        show_label=False, placeholder="Enter text and press ENTER", container=False
+    )
+with gr.Blocks(title='T-MoENet', theme=gr.themes.Default(), css=block_css) as demo:
+    gr.Markdown(title_markdown)
+    state = gr.State()
+    state_ = gr.State()
+    first_run = gr.State()
+    images_tensor = gr.State()
+
+    with gr.Row():
+        with gr.Column(scale=3):
+            video = gr.Video(label="Input Video")
+            cur_dir = os.path.dirname(os.path.abspath(__file__))
+            print(cur_dir)
+            gr.Examples(
+                examples=[
+                    [
+                        cur_dir + "/videos/3249402410.mp4",
+                        "what did the lady in black on the left do after she finished spreading the sauce on her pizza?",
+                    ],
+                    [
+                        cur_dir + "/videos/4882821564.mp4",
+                        "why did the boy clap his hands when he ran to the christmas tree?",
+                    ],
+                    [
+                        cur_dir + "/videos/6233408665.mp4",
+                        "what did the people on the sofa do after the lady in pink finished singing?",
+                    ],
+                ],
+                inputs=[video, textbox],
+            )
+
+        with gr.Column(scale=7):
+            chatbot = gr.Chatbot(label="T-MoENet", bubble_full_width=True)
+            with gr.Row():
+                with gr.Column(scale=2):
+                    textbox.render()
+                with gr.Column(scale=1, min_width=50):
+                    submit_btn = gr.Button(
+                        value="Send", variant="primary", interactive=True
+                    )
+
+    submit_btn.click(generate, [video, textbox, first_run, state, state_],
+                     [state, state_, chatbot, first_run, textbox,  video])
+
+demo.launch(share=True)

ckpts/deberta-v2-xlarge/config.json

@@ -0,0 +1,25 @@
+{
+	"model_type": "deberta-v2",
+	"attention_probs_dropout_prob": 0.1,
+	"hidden_act": "gelu",
+	"hidden_dropout_prob": 0.1,
+	"hidden_size": 1536,
+	"initializer_range": 0.02,
+	"intermediate_size": 6144,
+	"max_position_embeddings": 512,
+	"relative_attention": true,
+	"position_buckets": 256,
+	"norm_rel_ebd": "layer_norm",
+	"share_att_key": true,
+	"pos_att_type": "p2c|c2p",
+	"layer_norm_eps": 1e-7,
+	"conv_kernel_size": 3,
+	"conv_act": "gelu",
+	"max_relative_positions": -1,
+	"position_biased_input": false,
+	"num_attention_heads": 24,
+	"attention_head_size": 64,
+	"num_hidden_layers": 24,
+	"type_vocab_size": 0,
+	"vocab_size": 128100
+}

ckpts/deberta-v2-xlarge/pytorch_model.bin

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

ckpts/deberta-v2-xlarge/spm.model

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

ckpts/deberta-v2-xlarge/tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "do_lower_case": false,
+  "vocab_type": "spm"
+}

ckpts/model.pth

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

model/adapter.py

@@ -0,0 +1,77 @@
+import torch.nn as nn
+import torch
+import math
+
+class Adapter(nn.Module):
+    def __init__(
+        self, ds_factor, hidden_dim, ln_after=False, ln_before=False, dropout=0.1
+    ):
+        super().__init__()
+        assert not hidden_dim % ds_factor
+        self.down = nn.Linear(hidden_dim, hidden_dim // ds_factor)
+        self.act = nn.ReLU()
+        self.up = nn.Linear(hidden_dim // ds_factor, hidden_dim)
+        self.apply(self.init_weights)
+        self.ln_after = ln_after
+        self.ln_before = ln_before
+        self.dropout = dropout
+        if ln_after or ln_before:
+            self.ln = nn.LayerNorm(hidden_dim)
+        if dropout:
+            self.dropout = nn.Dropout(dropout)
+
+    def init_weights(self, m: nn.Module, std=1e-3):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.normal_(m.weight, std=std)
+            torch.nn.init.normal_(m.bias, std=std)
+            m.weight.data = torch.clamp(m.weight.data, min=-2 * std, max=2 * std)
+            m.bias.data = torch.clamp(m.bias.data, min=-2 * std, max=2 * std)
+        elif isinstance(m, nn.LayerNorm):
+            m.bias.data.zero_()
+            m.weight.data.fill_(1.0)
+
+    def forward(self, hidden_states):
+        if self.ln_before:
+            residual = self.ln(hidden_states)
+            residual = self.down(residual)
+        else:
+            residual = self.down(hidden_states)
+        residual = self.act(residual)
+        if self.dropout:
+            residual = self.dropout(residual)
+        residual = self.up(residual)
+        if self.ln_after:
+            residual = self.ln(hidden_states)
+        return hidden_states + residual
+
+
+class ST_Adapter(nn.Module):
+
+    def __init__(self, ds_factor, hidden_dim):
+        super().__init__()
+
+        self.down = nn.Linear(hidden_dim, hidden_dim // ds_factor)
+        self.conv = nn.Conv1d(
+            hidden_dim // ds_factor, hidden_dim // ds_factor,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+            groups=hidden_dim // ds_factor
+        )
+        self.up = nn.Linear(hidden_dim // ds_factor, hidden_dim)
+        nn.init.constant_(self.conv.weight, 0.)
+        nn.init.constant_(self.conv.bias, 0.)
+        nn.init.constant_(self.down.bias, 0.)
+        nn.init.constant_(self.up.bias, 0.)
+
+    def forward(self, x):
+        N, T, C = x.size()
+        ori_x = x
+        x = self.down(x)
+        x = x.permute(0, 2, 1).contiguous()
+        x = self.conv(x)
+        x = x.permute(0, 2, 1).contiguous()
+        x = self.up(x)
+        x = x + ori_x
+        return x
+    

model/deberta_moe.py

@@ -0,0 +1,1735 @@
+# coding=utf-8
+# Copyright 2020 Microsoft and the Hugging Face Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch DeBERTa-v2 model. """
+
+import math
+from collections.abc import Sequence
+from typing import Tuple, Optional
+
+import clip
+import numpy as np
+import torch
+from torch import _softmax_backward_data, nn
+from torch.nn import CrossEntropyLoss, LayerNorm
+
+from .adapter import Adapter
+from .moe import MoE
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import ModelOutput
+
+from transformers.modeling_utils import PreTrainedModel
+from transformers import DebertaV2Config, DebertaV2ForSequenceClassification
+from .evl import EVLTransformer, recursive_gumbel_softmax
+
+from transformers import pytorch_utils
+
+_CONFIG_FOR_DOC = "DebertaV2Config"
+_TOKENIZER_FOR_DOC = "DebertaV2Tokenizer"
+_CHECKPOINT_FOR_DOC = "microsoft/deberta-v2-xlarge"
+
+DEBERTA_V2_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "microsoft/deberta-v2-xlarge",
+    "microsoft/deberta-v2-xxlarge",
+    "microsoft/deberta-v2-xlarge-mnli",
+    "microsoft/deberta-v2-xxlarge-mnli",
+]
+
+class MaskedLMOutput(ModelOutput):
+    """
+    Base class for masked language models outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Masked language modeling (MLM) loss.
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    loss_moe: Optional[torch.FloatTensor] = None
+    loads: Optional[torch.FloatTensor] = None
+    embeddings: Optional[torch.FloatTensor] = None
+
+
+class BaseModelOutput(ModelOutput):
+    """
+    Base class for model's outputs, with potential hidden states and attentions.
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) of
+            shape `(batch_size, sequence_length, hidden_size)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    position_embeddings: torch.FloatTensor = None
+    attention_mask: torch.BoolTensor = None
+    loss_moe: torch.FloatTensor = None
+    video_g: torch.FloatTensor = None
+    loads: torch.LongTensor = None
+    embeddings: torch.FloatTensor = None
+
+
+# Copied from transformers.models.deberta.modeling_deberta.ContextPooler
+class ContextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.pooler_hidden_size, config.pooler_hidden_size)
+        self.dropout = StableDropout(config.pooler_dropout)
+        self.config = config
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+
+        context_token = hidden_states[:, 0]
+        context_token = self.dropout(context_token)
+        pooled_output = self.dense(context_token)
+        pooled_output = ACT2FN[self.config.pooler_hidden_act](pooled_output)
+        return pooled_output
+
+    @property
+    def output_dim(self):
+        return self.config.hidden_size
+
+
+# Copied from transformers.models.deberta.modeling_deberta.XSoftmax with deberta->deberta_v2
+class XSoftmax(torch.autograd.Function):
+    """
+    Masked Softmax which is optimized for saving memory
+
+    Args:
+        input (:obj:`torch.tensor`): The input tensor that will apply softmax.
+        mask (:obj:`torch.IntTensor`): The mask matrix where 0 indicate that element will be ignored in the softmax calculation.
+        dim (int): The dimension that will apply softmax
+
+    Example::
+
+          import torch
+          from transformers.models.deberta_v2.modeling_deberta_v2 import XSoftmax
+
+          # Make a tensor
+          x = torch.randn([4,20,100])
+
+          # Create a mask
+          mask = (x>0).int()
+
+          y = XSoftmax.apply(x, mask, dim=-1)
+    """
+
+    @staticmethod
+    def forward(self, input, mask, dim):
+        self.dim = dim
+        rmask = ~(mask.bool())
+
+        output = input.masked_fill(rmask, float("-inf"))
+        output = torch.softmax(output, self.dim)
+        output.masked_fill_(rmask, 0)
+        self.save_for_backward(output)
+        return output
+
+    @staticmethod
+    def backward(self, grad_output):
+        (output,) = self.saved_tensors
+        inputGrad = _softmax_backward_data(grad_output, output, self.dim, output.dtype)
+        return inputGrad, None, None
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DropoutContext
+class DropoutContext(object):
+    def __init__(self):
+        self.dropout = 0
+        self.mask = None
+        self.scale = 1
+        self.reuse_mask = True
+
+
+# Copied from transformers.models.deberta.modeling_deberta.get_mask
+def get_mask(input, local_context):
+    if not isinstance(local_context, DropoutContext):
+        dropout = local_context
+        mask = None
+    else:
+        dropout = local_context.dropout
+        dropout *= local_context.scale
+        mask = local_context.mask if local_context.reuse_mask else None
+
+    if dropout > 0 and mask is None:
+        mask = (1 - torch.empty_like(input).bernoulli_(1 - dropout)).bool()
+
+    if isinstance(local_context, DropoutContext):
+        if local_context.mask is None:
+            local_context.mask = mask
+
+    return mask, dropout
+
+
+# Copied from transformers.models.deberta.modeling_deberta.XDropout
+class XDropout(torch.autograd.Function):
+    """Optimized dropout function to save computation and memory by using mask operation instead of multiplication."""
+
+    @staticmethod
+    def forward(ctx, input, local_ctx):
+        mask, dropout = get_mask(input, local_ctx)
+        ctx.scale = 1.0 / (1 - dropout)
+        if dropout > 0:
+            ctx.save_for_backward(mask)
+            return input.masked_fill(mask, 0) * ctx.scale
+        else:
+            return input
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        if ctx.scale > 1:
+            (mask,) = ctx.saved_tensors
+            return grad_output.masked_fill(mask, 0) * ctx.scale, None
+        else:
+            return grad_output, None
+
+
+# Copied from transformers.models.deberta.modeling_deberta.StableDropout
+class StableDropout(nn.Module):
+    """
+    Optimized dropout module for stabilizing the training
+
+    Args:
+        drop_prob (float): the dropout probabilities
+    """
+
+    def __init__(self, drop_prob):
+        super().__init__()
+        self.drop_prob = drop_prob
+        self.count = 0
+        self.context_stack = None
+
+    def forward(self, x):
+        """
+        Call the module
+
+        Args:
+            x (:obj:`torch.tensor`): The input tensor to apply dropout
+        """
+        if self.training and self.drop_prob > 0:
+            return XDropout.apply(x, self.get_context())
+        return x
+
+    def clear_context(self):
+        self.count = 0
+        self.context_stack = None
+
+    def init_context(self, reuse_mask=True, scale=1):
+        if self.context_stack is None:
+            self.context_stack = []
+        self.count = 0
+        for c in self.context_stack:
+            c.reuse_mask = reuse_mask
+            c.scale = scale
+
+    def get_context(self):
+        if self.context_stack is not None:
+            if self.count >= len(self.context_stack):
+                self.context_stack.append(DropoutContext())
+            ctx = self.context_stack[self.count]
+            ctx.dropout = self.drop_prob
+            self.count += 1
+            return ctx
+        else:
+            return self.drop_prob
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaSelfOutput with DebertaLayerNorm->LayerNorm
+class DebertaV2SelfOutput(nn.Module):
+    def __init__(self, config, ds_factor, dropout, add_moe, gating):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.add_moe = add_moe
+        if not self.add_moe and ds_factor:
+            self.adapter = Adapter(ds_factor, config.hidden_size, dropout=dropout)
+        else:
+            self.moe_layer = MoE(ds_factor = ds_factor, moe_input_size=config.hidden_size, dropout=dropout, num_experts=4, top_k=2, gating=gating)
+
+    def forward(self, hidden_states, input_tensor, temporal_factor = None, train_mode = True):
+        hidden_states = self.dense(hidden_states)
+        if not self.add_moe:
+            hidden_states = self.adapter(hidden_states)
+        else:
+            hidden_states, loss_moe, load = self.moe_layer(temporal_factor, hidden_states, train=train_mode)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        
+        if not self.add_moe:
+            return hidden_states, None, None
+        
+        return hidden_states, loss_moe, load
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaAttention with Deberta->DebertaV2
+class DebertaV2Attention(nn.Module):
+    def __init__(self, config, ds_factor, dropout, add_moe = False, gating='linear'):
+        super().__init__()
+        self.self = DisentangledSelfAttention(config)
+        self.output = DebertaV2SelfOutput(config, ds_factor, dropout, add_moe, gating)
+        self.config = config
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        temporal_factor=None,
+        train_mode=True
+    ):
+        self_output = self.self(
+            hidden_states,
+            attention_mask,
+            return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+        )
+        if return_att:
+            self_output, att_matrix = self_output
+        if query_states is None:
+            query_states = hidden_states
+        attention_output, loss_moe, load = self.output(self_output, query_states, temporal_factor, train_mode)
+
+        if return_att:
+            return (attention_output, att_matrix, loss_moe)
+        else:
+            return attention_output, loss_moe, load
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->DebertaV2
+class DebertaV2Intermediate(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
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaOutput with DebertaLayerNorm->LayerNorm
+class DebertaV2Output(nn.Module):
+    def __init__(self, config, ds_factor, dropout, add_moe = False, gating='linear',layer_id=0):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+        self.ds_factor = ds_factor
+        self.add_moe = add_moe
+        if not self.add_moe and self.ds_factor:
+            self.adapter = Adapter(ds_factor, config.hidden_size, dropout=dropout)
+        elif self.add_moe:
+            self.moe_layer = MoE(ds_factor=ds_factor, moe_input_size=config.hidden_size, dropout=dropout, num_experts=4, top_k=1, gating=gating, layer_id=layer_id)
+            #self.adapter = Adapter(ds_factor, config.hidden_size, dropout=dropout)
+
+    def forward(self, hidden_states, input_tensor, temporal_factor, train_mode):
+        hidden_states = self.dense(hidden_states)
+        if not self.add_moe and self.ds_factor:
+            hidden_states = self.adapter(hidden_states)
+        elif self.add_moe:
+            hidden_states, loss_moe, load = self.moe_layer(temporal_factor, hidden_states, train=train_mode)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+
+        if not self.add_moe:
+            return hidden_states, None, None
+        
+        return hidden_states, loss_moe, load
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaLayer with Deberta->DebertaV2
+class DebertaV2Layer(nn.Module):
+    def __init__(
+        self,
+        config,
+        ds_factor_attn,
+        ds_factor_ff,
+        dropout,
+        layer_id,
+    ):
+        super().__init__()
+        self.layer_id = layer_id
+        self.add_moe = False
+        
+        #if layer_id >= config.num_hidden_layers - 2:
+        #    self.add_moe = True
+        
+        if layer_id < 2:
+           self.add_moe = True
+
+        self.attention = DebertaV2Attention(config, ds_factor_attn, dropout, False)
+        self.intermediate = DebertaV2Intermediate(config)
+        self.output = DebertaV2Output(config, ds_factor_ff, dropout, self.add_moe, gating="linear", layer_id = layer_id)
+
+    def forward(
+        self,
+        temporal_factor,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+        train_mode=True,
+    ):
+        attention_output = self.attention(
+            hidden_states,
+            attention_mask,
+            return_att=return_att,
+            query_states=query_states,
+            relative_pos=relative_pos,
+            rel_embeddings=rel_embeddings,
+            temporal_factor=temporal_factor,
+            train_mode=train_mode
+        )
+
+        if return_att:
+            attention_output, att_matrix, loss_moe_attn = attention_output
+        else:
+            attention_output, loss_moe_attn, load = attention_output
+        intermediate_output = self.intermediate(attention_output)
+        layer_output, loss_moe_ffn, load = self.output(intermediate_output, attention_output, temporal_factor=temporal_factor, train_mode=train_mode)
+        
+        loss_moe = loss_moe_attn if loss_moe_attn else loss_moe_ffn
+        if return_att:
+            return (layer_output, att_matrix)
+        
+        
+        return layer_output, loss_moe, load
+
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        kernel_size = getattr(config, "conv_kernel_size", 3)
+        groups = getattr(config, "conv_groups", 1)
+        self.conv_act = getattr(config, "conv_act", "tanh")
+        self.conv = nn.Conv1d(
+            config.hidden_size,
+            config.hidden_size,
+            kernel_size,
+            padding=(kernel_size - 1) // 2,
+            groups=groups,
+        )
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+    def forward(self, hidden_states, residual_states, input_mask):
+        out = (
+            self.conv(hidden_states.permute(0, 2, 1).contiguous())
+            .permute(0, 2, 1)
+            .contiguous()
+        )
+        rmask = (1 - input_mask).bool()
+        out.masked_fill_(rmask.unsqueeze(-1).expand(out.size()), 0)
+        out = ACT2FN[self.conv_act](self.dropout(out))
+
+        layer_norm_input = residual_states + out
+        output = self.LayerNorm(layer_norm_input).to(layer_norm_input)
+
+        if input_mask is None:
+            output_states = output
+        else:
+            if input_mask.dim() != layer_norm_input.dim():
+                if input_mask.dim() == 4:
+                    input_mask = input_mask.squeeze(1).squeeze(1)
+                input_mask = input_mask.unsqueeze(2)
+
+            input_mask = input_mask.to(output.dtype)
+            output_states = output * input_mask
+
+        return output_states
+
+
+class DebertaV2Encoder(nn.Module):
+    """Modified BertEncoder with relative position bias support"""
+
+    def __init__(
+        self,
+        config,
+        ds_factor_attn,
+        ds_factor_ff,
+        dropout,
+    ):
+        super().__init__()
+
+        self.layer = nn.ModuleList(
+            [
+                DebertaV2Layer(
+                    config,
+                    ds_factor_attn,
+                    ds_factor_ff,
+                    dropout,
+                    _,
+                )
+                for _ in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.relative_attention = getattr(config, "relative_attention", False)
+
+        if self.relative_attention:
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+
+            self.position_buckets = getattr(config, "position_buckets", -1)
+            pos_ebd_size = self.max_relative_positions * 2
+
+            if self.position_buckets > 0:
+                pos_ebd_size = self.position_buckets * 2
+
+            self.rel_embeddings = nn.Embedding(pos_ebd_size, config.hidden_size)
+
+        self.norm_rel_ebd = [
+            x.strip()
+            for x in getattr(config, "norm_rel_ebd", "none").lower().split("|")
+        ]
+
+        if "layer_norm" in self.norm_rel_ebd:
+            self.LayerNorm = LayerNorm(
+                config.hidden_size, config.layer_norm_eps, elementwise_affine=True
+            )
+
+        self.conv = (
+            ConvLayer(config) if getattr(config, "conv_kernel_size", 0) > 0 else None
+        )
+
+    def get_rel_embedding(self):
+        rel_embeddings = self.rel_embeddings.weight if self.relative_attention else None
+        if rel_embeddings is not None and ("layer_norm" in self.norm_rel_ebd):
+            rel_embeddings = self.LayerNorm(rel_embeddings)
+        return rel_embeddings
+
+    def get_attention_mask(self, attention_mask):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = extended_attention_mask * extended_attention_mask.squeeze(
+                -2
+            ).unsqueeze(-1)
+            attention_mask = attention_mask.byte()
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask.unsqueeze(1)
+
+        return attention_mask
+
+    def get_rel_pos(self, hidden_states, query_states=None, relative_pos=None):
+        if self.relative_attention and relative_pos is None:
+            q = (
+                query_states.size(-2)
+                if query_states is not None
+                else hidden_states.size(-2)
+            )
+            relative_pos = build_relative_position(
+                q,
+                hidden_states.size(-2),
+                bucket_size=self.position_buckets,
+                max_position=self.max_relative_positions,
+            )
+        return relative_pos
+
+    def forward(
+        self,
+        temporal_factor,
+        hidden_states,
+        attention_mask,
+        output_hidden_states=True,
+        output_attentions=False,
+        query_states=None,
+        relative_pos=None,
+        return_dict=True,
+        train_mode=True
+    ):
+        if attention_mask.dim() <= 2:
+            input_mask = attention_mask
+        else:
+            input_mask = (attention_mask.sum(-2) > 0).byte()
+        attention_mask = self.get_attention_mask(attention_mask)
+        relative_pos = self.get_rel_pos(hidden_states, query_states, relative_pos)
+
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        if isinstance(hidden_states, Sequence):
+            next_kv = hidden_states[0]
+        else:
+            next_kv = hidden_states
+        rel_embeddings = self.get_rel_embedding()
+        output_states = next_kv
+
+        loss_moe = 0
+        loads = []
+        embeddings = []
+
+        for i, layer_module in enumerate(self.layer):
+
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (output_states,)
+
+            output_states, _, load = layer_module(
+                temporal_factor,
+                next_kv,
+                attention_mask,
+                output_attentions,
+                query_states=query_states,
+                relative_pos=relative_pos,
+                rel_embeddings=rel_embeddings,
+                train_mode=train_mode
+            )
+            if isinstance(load, torch.Tensor):
+                loads.append(load)
+
+            if _:
+                loss_moe = loss_moe + _
+
+            if output_attentions:
+                output_states, att_m = output_states
+
+            if i == 0 and self.conv is not None:
+                output_states = self.conv(hidden_states, output_states, input_mask)
+
+            if query_states is not None:
+                query_states = output_states
+                if isinstance(hidden_states, Sequence):
+                    next_kv = hidden_states[i + 1] if i + 1 < len(self.layer) else None
+            else:
+                next_kv = output_states
+
+            if output_attentions:
+                all_attentions = all_attentions + (att_m,)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (output_states,)
+
+        if not return_dict:
+            return tuple(
+                v
+                for v in [output_states, all_hidden_states, all_attentions]
+                if v is not None
+            )
+        
+        if len(loads)>0:
+            loads = torch.stack(loads, dim = 0)
+
+        if len(embeddings) >0:
+            embeddings = torch.cat(embeddings, dim=0)
+
+        return BaseModelOutput(
+            last_hidden_state=output_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            loss_moe=loss_moe,
+            loads=loads
+        )
+
+
+def make_log_bucket_position(relative_pos, bucket_size, max_position):
+    sign = np.sign(relative_pos)
+    mid = bucket_size // 2
+    abs_pos = np.where(
+        (relative_pos < mid) & (relative_pos > -mid), mid - 1, np.abs(relative_pos)
+    )
+    log_pos = (
+        np.ceil(np.log(abs_pos / mid) / np.log((max_position - 1) / mid) * (mid - 1))
+        + mid
+    )
+    bucket_pos = np.where(abs_pos <= mid, relative_pos, log_pos * sign).astype(np.int)
+    return bucket_pos
+
+
+def build_relative_position(query_size, key_size, bucket_size=-1, max_position=-1):
+    """
+    Build relative position according to the query and key
+
+    We assume the absolute position of query :math:`P_q` is range from (0, query_size) and the absolute position of key
+    :math:`P_k` is range from (0, key_size), The relative positions from query to key is :math:`R_{q \\rightarrow k} =
+    P_q - P_k`
+
+    Args:
+        query_size (int): the length of query
+        key_size (int): the length of key
+        bucket_size (int): the size of position bucket
+        max_position (int): the maximum allowed absolute position
+
+    Return:
+        :obj:`torch.LongTensor`: A tensor with shape [1, query_size, key_size]
+
+    """
+    q_ids = np.arange(0, query_size)
+    k_ids = np.arange(0, key_size)
+    rel_pos_ids = q_ids[:, None] - np.tile(k_ids, (q_ids.shape[0], 1))
+    if bucket_size > 0 and max_position > 0:
+        rel_pos_ids = make_log_bucket_position(rel_pos_ids, bucket_size, max_position)
+    rel_pos_ids = torch.tensor(rel_pos_ids, dtype=torch.long)
+    rel_pos_ids = rel_pos_ids[:query_size, :]
+    rel_pos_ids = rel_pos_ids.unsqueeze(0)
+    return rel_pos_ids
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.c2p_dynamic_expand
+def c2p_dynamic_expand(c2p_pos, query_layer, relative_pos):
+    return c2p_pos.expand(
+        [
+            query_layer.size(0),
+            query_layer.size(1),
+            query_layer.size(2),
+            relative_pos.size(-1),
+        ]
+    )
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.p2c_dynamic_expand
+def p2c_dynamic_expand(c2p_pos, query_layer, key_layer):
+    return c2p_pos.expand(
+        [
+            query_layer.size(0),
+            query_layer.size(1),
+            key_layer.size(-2),
+            key_layer.size(-2),
+        ]
+    )
+
+
+@torch.jit.script
+# Copied from transformers.models.deberta.modeling_deberta.pos_dynamic_expand
+def pos_dynamic_expand(pos_index, p2c_att, key_layer):
+    return pos_index.expand(
+        p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2))
+    )
+
+
+class DisentangledSelfAttention(nn.Module):
+    """
+    Disentangled self-attention module
+
+    Parameters:
+        config (:obj:`DebertaV2Config`):
+            A model config class instance with the configuration to build a new model. The schema is similar to
+            `BertConfig`, for more details, please refer :class:`~transformers.DebertaV2Config`
+
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+                f"heads ({config.num_attention_heads})"
+            )
+        self.num_attention_heads = config.num_attention_heads
+        _attention_head_size = config.hidden_size // config.num_attention_heads
+        self.attention_head_size = getattr(
+            config, "attention_head_size", _attention_head_size
+        )
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.query_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.key_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+        self.value_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
+
+        self.share_att_key = getattr(config, "share_att_key", False)
+        self.pos_att_type = (
+            config.pos_att_type if config.pos_att_type is not None else []
+        )
+        self.relative_attention = getattr(config, "relative_attention", False)
+
+        if self.relative_attention:
+            self.position_buckets = getattr(config, "position_buckets", -1)
+            self.max_relative_positions = getattr(config, "max_relative_positions", -1)
+            if self.max_relative_positions < 1:
+                self.max_relative_positions = config.max_position_embeddings
+            self.pos_ebd_size = self.max_relative_positions
+            if self.position_buckets > 0:
+                self.pos_ebd_size = self.position_buckets
+
+            self.pos_dropout = StableDropout(config.hidden_dropout_prob)
+
+            if not self.share_att_key:
+                if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                    self.pos_key_proj = nn.Linear(
+                        config.hidden_size, self.all_head_size, bias=True
+                    )
+                if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                    self.pos_query_proj = nn.Linear(
+                        config.hidden_size, self.all_head_size
+                    )
+
+        self.dropout = StableDropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x, attention_heads):
+        new_x_shape = x.size()[:-1] + (attention_heads, -1)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3).contiguous().view(-1, x.size(1), x.size(-1))
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        return_att=False,
+        query_states=None,
+        relative_pos=None,
+        rel_embeddings=None,
+    ):
+        """
+        Call the module
+
+        Args:
+            hidden_states (:obj:`torch.FloatTensor`):
+                Input states to the module usually the output from previous layer, it will be the Q,K and V in
+                `Attention(Q,K,V)`
+
+            attention_mask (:obj:`torch.ByteTensor`):
+                An attention mask matrix of shape [`B`, `N`, `N`] where `B` is the batch size, `N` is the maximum
+                sequence length in which element [i,j] = `1` means the `i` th token in the input can attend to the `j`
+                th token.
+
+            return_att (:obj:`bool`, optional):
+                Whether return the attention matrix.
+
+            query_states (:obj:`torch.FloatTensor`, optional):
+                The `Q` state in `Attention(Q,K,V)`.
+
+            relative_pos (:obj:`torch.LongTensor`):
+                The relative position encoding between the tokens in the sequence. It's of shape [`B`, `N`, `N`] with
+                values ranging in [`-max_relative_positions`, `max_relative_positions`].
+
+            rel_embeddings (:obj:`torch.FloatTensor`):
+                The embedding of relative distances. It's a tensor of shape [:math:`2 \\times
+                \\text{max_relative_positions}`, `hidden_size`].
+
+
+        """
+        if query_states is None:
+            query_states = hidden_states
+        query_layer = self.transpose_for_scores(
+            self.query_proj(query_states), self.num_attention_heads
+        )
+        key_layer = self.transpose_for_scores(
+            self.key_proj(hidden_states), self.num_attention_heads
+        )
+        value_layer = self.transpose_for_scores(
+            self.value_proj(hidden_states), self.num_attention_heads
+        )
+
+        rel_att = None
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        scale_factor = 1
+        if "c2p" in self.pos_att_type:
+            scale_factor += 1
+        if "p2c" in self.pos_att_type:
+            scale_factor += 1
+        if "p2p" in self.pos_att_type:
+            scale_factor += 1
+        scale = math.sqrt(query_layer.size(-1) * scale_factor)
+        attention_scores = torch.bmm(query_layer, key_layer.transpose(-1, -2)) / scale
+        if self.relative_attention:
+            rel_embeddings = self.pos_dropout(rel_embeddings)
+            rel_att = self.disentangled_attention_bias(
+                query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+            )
+
+        if rel_att is not None:
+            attention_scores = attention_scores + rel_att
+        attention_scores = attention_scores
+        attention_scores = attention_scores.view(
+            -1,
+            self.num_attention_heads,
+            attention_scores.size(-2),
+            attention_scores.size(-1),
+        )
+
+        # bsz x height x length x dimension
+        attention_probs = XSoftmax.apply(attention_scores, attention_mask, -1)
+        attention_probs = self.dropout(attention_probs)
+        context_layer = torch.bmm(
+            attention_probs.view(
+                -1, attention_probs.size(-2), attention_probs.size(-1)
+            ),
+            value_layer,
+        )
+        context_layer = (
+            context_layer.view(
+                -1,
+                self.num_attention_heads,
+                context_layer.size(-2),
+                context_layer.size(-1),
+            )
+            .permute(0, 2, 1, 3)
+            .contiguous()
+        )
+        new_context_layer_shape = context_layer.size()[:-2] + (-1,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+        if return_att:
+            return (context_layer, attention_probs)
+        else:
+            return context_layer
+
+    def disentangled_attention_bias(
+        self, query_layer, key_layer, relative_pos, rel_embeddings, scale_factor
+    ):
+        if relative_pos is None:
+            q = query_layer.size(-2)
+            relative_pos = build_relative_position(
+                q,
+                key_layer.size(-2),
+                bucket_size=self.position_buckets,
+                max_position=self.max_relative_positions,
+            )
+        if relative_pos.dim() == 2:
+            relative_pos = relative_pos.unsqueeze(0).unsqueeze(0)
+        elif relative_pos.dim() == 3:
+            relative_pos = relative_pos.unsqueeze(1)
+        # bsz x height x query x key
+        elif relative_pos.dim() != 4:
+            raise ValueError(
+                f"Relative position ids must be of dim 2 or 3 or 4. {relative_pos.dim()}"
+            )
+
+        att_span = self.pos_ebd_size
+        relative_pos = relative_pos.long().to(query_layer.device)
+
+        rel_embeddings = rel_embeddings[
+            self.pos_ebd_size - att_span : self.pos_ebd_size + att_span, :
+        ].unsqueeze(0)
+        if self.share_att_key:
+            pos_query_layer = self.transpose_for_scores(
+                self.query_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
+            pos_key_layer = self.transpose_for_scores(
+                self.key_proj(rel_embeddings), self.num_attention_heads
+            ).repeat(query_layer.size(0) // self.num_attention_heads, 1, 1)
+        else:
+            if "c2p" in self.pos_att_type or "p2p" in self.pos_att_type:
+                pos_key_layer = self.transpose_for_scores(
+                    self.pos_key_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(
+                    query_layer.size(0) // self.num_attention_heads, 1, 1
+                )  # .split(self.all_head_size, dim=-1)
+            if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+                pos_query_layer = self.transpose_for_scores(
+                    self.pos_query_proj(rel_embeddings), self.num_attention_heads
+                ).repeat(
+                    query_layer.size(0) // self.num_attention_heads, 1, 1
+                )  # .split(self.all_head_size, dim=-1)
+
+        score = 0
+        # content->position
+        if "c2p" in self.pos_att_type:
+            scale = math.sqrt(pos_key_layer.size(-1) * scale_factor)
+            c2p_att = torch.bmm(query_layer, pos_key_layer.transpose(-1, -2))
+            c2p_pos = torch.clamp(relative_pos + att_span, 0, att_span * 2 - 1)
+            c2p_att = torch.gather(
+                c2p_att,
+                dim=-1,
+                index=c2p_pos.squeeze(0).expand(
+                    [query_layer.size(0), query_layer.size(1), relative_pos.size(-1)]
+                ),
+            )
+            score += c2p_att / scale
+
+        # position->content
+        if "p2c" in self.pos_att_type or "p2p" in self.pos_att_type:
+            scale = math.sqrt(pos_query_layer.size(-1) * scale_factor)
+            if key_layer.size(-2) != query_layer.size(-2):
+                r_pos = build_relative_position(
+                    key_layer.size(-2),
+                    key_layer.size(-2),
+                    bucket_size=self.position_buckets,
+                    max_position=self.max_relative_positions,
+                ).to(query_layer.device)
+                r_pos = r_pos.unsqueeze(0)
+            else:
+                r_pos = relative_pos
+
+            p2c_pos = torch.clamp(-r_pos + att_span, 0, att_span * 2 - 1)
+            if query_layer.size(-2) != key_layer.size(-2):
+                pos_index = relative_pos[:, :, :, 0].unsqueeze(-1)
+
+        if "p2c" in self.pos_att_type:
+            p2c_att = torch.bmm(key_layer, pos_query_layer.transpose(-1, -2))
+            p2c_att = torch.gather(
+                p2c_att,
+                dim=-1,
+                index=p2c_pos.squeeze(0).expand(
+                    [query_layer.size(0), key_layer.size(-2), key_layer.size(-2)]
+                ),
+            ).transpose(-1, -2)
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2c_att = torch.gather(
+                    p2c_att,
+                    dim=-2,
+                    index=pos_index.expand(
+                        p2c_att.size()[:2] + (pos_index.size(-2), key_layer.size(-2))
+                    ),
+                )
+            score += p2c_att / scale
+
+        # position->position
+        if "p2p" in self.pos_att_type:
+            pos_query = pos_query_layer[:, :, att_span:, :]
+            p2p_att = torch.matmul(pos_query, pos_key_layer.transpose(-1, -2))
+            p2p_att = p2p_att.expand(query_layer.size()[:2] + p2p_att.size()[2:])
+            if query_layer.size(-2) != key_layer.size(-2):
+                p2p_att = torch.gather(
+                    p2p_att,
+                    dim=-2,
+                    index=pos_index.expand(
+                        query_layer.size()[:2] + (pos_index.size(-2), p2p_att.size(-1))
+                    ),
+                )
+            p2p_att = torch.gather(
+                p2p_att,
+                dim=-1,
+                index=c2p_pos.expand(
+                    [
+                        query_layer.size(0),
+                        query_layer.size(1),
+                        query_layer.size(2),
+                        relative_pos.size(-1),
+                    ]
+                ),
+            )
+            score += p2p_att
+
+        return score
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaEmbeddings with DebertaLayerNorm->LayerNorm
+class DebertaV2Embeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(
+        self,
+        config,
+        features_dim,
+        add_video_feat=False,
+        max_feats = 10
+    ):
+        super().__init__()
+        pad_token_id = getattr(config, "pad_token_id", 0)
+        self.embedding_size = getattr(config, "embedding_size", config.hidden_size)
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, self.embedding_size, padding_idx=pad_token_id
+        )
+
+        self.position_biased_input = getattr(config, "position_biased_input", True)
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, self.embedding_size
+        )  # it is used for the decoder anyway
+
+        if config.type_vocab_size > 0:
+            self.token_type_embeddings = nn.Embedding(
+                config.type_vocab_size, self.embedding_size
+            )
+
+        if self.embedding_size != config.hidden_size:
+            self.embed_proj = nn.Linear(
+                self.embedding_size, config.hidden_size, bias=False
+            )
+        self.LayerNorm = LayerNorm(config.hidden_size, config.layer_norm_eps)
+        self.dropout = StableDropout(config.hidden_dropout_prob)
+        self.config = config
+
+        # 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.add_video_feat = add_video_feat
+        self.features_dim = features_dim
+        if self.features_dim:
+            self.linear_video = nn.Linear(features_dim, config.hidden_size)
+            if self.add_video_feat:
+                self.evl = EVLTransformer(max_feats, decoder_num_layers=1, 
+                                          decoder_qkv_dim=768, add_video_feat=self.add_video_feat,
+                                          add_mask=True)
+            #self.evl = ConvNet()
+
+    def get_video_embedding(self, video, video_mask):
+        
+        if self.add_video_feat:
+            video_g = self.evl(video, video_mask)
+            video_feat = self.linear_video(video)
+            video_feat_l = torch.cat([video_g, video_feat], dim = 1)
+            
+        else:
+            video_feat_l = self.linear_video(video)
+            video_feat_tmp = video_feat_l * video_mask.unsqueeze(-1)
+            video_g = torch.sum(video_feat_tmp, dim = 1) / video_mask.sum(dim = 1, keepdim=True)
+        return video_g, video_feat_l
+
+    def forward(
+        self,
+        input_ids=None,
+        token_type_ids=None,
+        position_ids=None,
+        mask=None,
+        inputs_embeds=None,
+        video=None,
+        video_mask=None
+    ):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+            if self.features_dim and video is not None:
+                video_global, video = self.get_video_embedding(video, video_mask)
+                inputs_embeds = torch.cat([video, inputs_embeds], 1)
+                input_shape = inputs_embeds[:, :, 0].shape
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(
+                input_shape, dtype=torch.long, device=self.position_ids.device
+            )
+
+        if self.position_embeddings is not None:
+            position_embeddings = self.position_embeddings(position_ids.long())
+        else:
+            position_embeddings = torch.zeros_like(inputs_embeds)
+
+        embeddings = inputs_embeds
+        if self.position_biased_input:
+            embeddings = embeddings + position_embeddings
+        if self.config.type_vocab_size > 0:
+            token_type_embeddings = self.token_type_embeddings(token_type_ids)
+            embeddings = embeddings + token_type_embeddings
+
+        if self.embedding_size != self.config.hidden_size:
+            embeddings = self.embed_proj(embeddings)
+
+        embeddings = self.LayerNorm(embeddings)
+
+        if mask is not None:
+            if mask.dim() != embeddings.dim():
+                if mask.dim() == 4:
+                    mask = mask.squeeze(1).squeeze(1)
+                mask = mask.unsqueeze(2)
+            mask = mask.to(embeddings.dtype)
+
+            embeddings = embeddings * mask
+
+        embeddings = self.dropout(embeddings)
+        return {
+            "embeddings": embeddings,
+            "position_embeddings": position_embeddings,
+            "video_global": video_global
+        }
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaPreTrainedModel with Deberta->DebertaV2
+
+
+class DebertaV2PreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = DebertaV2Config
+    base_model_prefix = "deberta"
+    _keys_to_ignore_on_load_missing = ["position_ids"]
+    _keys_to_ignore_on_load_unexpected = ["position_embeddings"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self._register_load_state_dict_pre_hook(self._pre_load_hook)
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            # 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)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _pre_load_hook(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        """
+        Removes the classifier if it doesn't have the correct number of labels.
+        """
+        self_state = self.state_dict()
+        if (
+            ("classifier.weight" in self_state)
+            and ("classifier.weight" in state_dict)
+            and self_state["classifier.weight"].size()
+            != state_dict["classifier.weight"].size()
+        ):
+            print(
+                f"The checkpoint classifier head has a shape {state_dict['classifier.weight'].size()} and this model "
+                f"classifier head has a shape {self_state['classifier.weight'].size()}. Ignoring the checkpoint "
+                f"weights. You should train your model on new data."
+            )
+            del state_dict["classifier.weight"]
+            if "classifier.bias" in state_dict:
+                del state_dict["classifier.bias"]
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaModel with Deberta->DebertaV2
+class DebertaV2Model(DebertaV2PreTrainedModel):
+    def __init__(
+        self,
+        config,
+        max_feats=10,
+        features_dim=768,
+        freeze_lm=False,
+        ds_factor_attn=8,
+        ds_factor_ff=8,
+        ft_ln=False,
+        dropout=0.1,
+        add_video_feat = False,
+        freeze_ad=False,
+    ):
+        super().__init__(config)
+        
+        self.embeddings = DebertaV2Embeddings(
+            config,
+            features_dim,
+            add_video_feat,
+            max_feats
+        )
+        self.encoder = DebertaV2Encoder(
+            config,
+            ds_factor_attn,
+            ds_factor_ff,
+            dropout,
+        )
+        self.z_steps = 0
+        self.config = config
+
+        self.features_dim = features_dim
+        self.max_feats = max_feats
+        if freeze_lm:
+            for n, p in self.named_parameters():
+                #if (not "linear_video" in n) and (not "adapter" in n):
+                #    if ft_ln and "LayerNorm" in n:
+                #        continue
+                #    else:
+                #        p.requires_grad_(False)
+                if not freeze_ad:
+                    if (not "evl" in n) and (not "linear_video" in n) and (not "adapter" in n) and (not "moe" in n):
+                        if ft_ln and "LayerNorm" in n:
+                            continue
+                        else:
+                            p.requires_grad_(False)
+                
+                else:
+                    if not "evl" in n:
+                        p.requires_grad_(False)
+                
+
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings.word_embeddings = new_embeddings
+
+    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
+        """
+        raise NotImplementedError(
+            "The prune function is not implemented in DeBERTa model."
+        )
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        video=None,
+        video_mask=None,
+        train_mode = True
+    ):
+        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
+        )
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_ids and inputs_embeds at the same time"
+            )
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+
+        if self.features_dim and video is not None:
+            if video_mask is None:
+                video_shape = video[:, :, 0].size()
+                video_mask = torch.ones(video_shape, device=device)
+            attention_mask = torch.cat([video_mask, attention_mask], 1)
+            input_shape = attention_mask.size()
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            video=video,
+            video_mask=video_mask[:, 1:] if video_mask.shape[1] != video.shape[1] else video_mask
+        )
+        embedding_output, position_embeddings, video_g = (
+            embedding_output["embeddings"],
+            embedding_output["position_embeddings"],
+            embedding_output["video_global"]
+        )
+
+        video_g = video_g.squeeze()
+        encoder_outputs = self.encoder(
+            video_g,
+            embedding_output,
+            attention_mask,
+            output_hidden_states=True,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+            train_mode=train_mode
+        )
+        encoded_layers = encoder_outputs[1]
+        loss_moe =encoder_outputs.loss_moe
+
+        if self.z_steps > 1:
+            hidden_states = encoded_layers[-2]
+            layers = [self.encoder.layer[-1] for _ in range(self.z_steps)]
+            query_states = encoded_layers[-1]
+            rel_embeddings = self.encoder.get_rel_embedding()
+            attention_mask = self.encoder.get_attention_mask(attention_mask)
+            rel_pos = self.encoder.get_rel_pos(embedding_output)
+            for layer in layers[1:]:
+                query_states = layer(
+                    hidden_states,
+                    attention_mask,
+                    return_att=False,
+                    query_states=query_states,
+                    relative_pos=rel_pos,
+                    rel_embeddings=rel_embeddings,
+                )
+                encoded_layers.append(query_states)
+
+        sequence_output = encoded_layers[-1]
+
+        if not return_dict:
+            return (sequence_output,) + encoder_outputs[
+                (1 if output_hidden_states else 2) :
+            ]
+
+        return BaseModelOutput(
+            last_hidden_state=sequence_output,
+            hidden_states=encoder_outputs.hidden_states
+            if output_hidden_states
+            else None,
+            attentions=encoder_outputs.attentions,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            video_g=video_g,
+            loss_moe = loss_moe,
+            loads=encoder_outputs.loads
+        )
+
+
+# Copied from transformers.models.deberta.modeling_deberta.DebertaForMaskedLM with Deberta->DebertaV2
+class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
+    _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,
+        max_feats=10,
+        features_dim=768,
+        freeze_lm=True,
+        freeze_mlm=True,
+        ds_factor_attn=8,
+        ds_factor_ff=8,
+        ft_ln=True,
+        dropout=0.1,
+        n_ans=0,
+        freeze_last=True,
+        add_video_feat = False,
+        freeze_ad=False,
+        add_temporal_trans = False
+    ):
+        """
+        :param config: BiLM configuration
+        :param max_feats: maximum number of frames used by the model
+        :param features_dim: embedding dimension of the visual features, set = 0 for text-only mode
+        :param freeze_lm: whether to freeze or not the language model (Transformer encoder + token embedder)
+        :param freeze_mlm: whether to freeze or not the MLM head
+        :param ds_factor_attn: downsampling factor for the adapter after self-attention, no adapter if set to 0
+        :param ds_factor_ff: downsampling factor for the adapter after feed-forward, no adapter if set to 0
+        :param ft_ln: whether to finetune or not the normalization layers
+        :param dropout: dropout probability in the adapter
+        :param n_ans: number of answers in the downstream vocabulary, set = 0 during cross-modal training
+        :param freeze_last: whether to freeze or not the answer embedding module
+        """
+        super().__init__(config)
+
+        # self.clip, _ = clip.load("ViT-L/14")
+        # for p in self.clip.parameters():
+        #    p.requires_grad_(False)
+
+        self.deberta = DebertaV2Model(
+            config,
+            max_feats,
+            features_dim,
+            freeze_lm,
+            ds_factor_attn,
+            ds_factor_ff,
+            ft_ln,
+            dropout,
+            add_video_feat,
+            freeze_ad
+        )
+
+        self.add_video_feat = add_video_feat
+        self.lm_predictions = DebertaV2OnlyMLMHead(config)
+        self.features_dim = features_dim
+        if freeze_mlm:
+            for n, p in self.lm_predictions.named_parameters():
+                if ft_ln and "LayerNorm" in n:
+                    continue
+                else:
+                    p.requires_grad_(False)
+
+        self.init_weights()
+        self.n_ans = n_ans
+        if n_ans:
+            self.answer_embeddings = nn.Embedding(
+                n_ans, self.deberta.embeddings.embedding_size
+            )
+            self.answer_bias = nn.Parameter(torch.zeros(n_ans))
+            if freeze_last:
+                self.answer_embeddings.requires_grad_(False)
+                self.answer_bias.requires_grad_(False)
+
+    def set_answer_embeddings(self, a2tok, freeze_last=True):
+        a2v = self.deberta.embeddings.word_embeddings(a2tok)  # answer embeddings   (ans_vocab_num, 1, dim)
+        pad_token_id = getattr(self.config, "pad_token_id", 0)
+        sum_tokens = (a2tok != pad_token_id).sum(1, keepdims=True)  # n_ans  (1000, 1) n_tokens
+        if len(a2v) != self.n_ans:  # reinitialize the answer embeddings
+            assert not self.training
+            self.n_ans = len(a2v)
+            self.answer_embeddings = nn.Embedding(
+                self.n_ans, self.deberta.embeddings.embedding_size
+            ).to(self.device)
+            self.answer_bias.requires_grad = False
+            self.answer_bias.resize_(self.n_ans)
+        self.answer_embeddings.weight.data = torch.div(
+            (a2v * (a2tok != pad_token_id).float()[:, :, None]).sum(1),
+            sum_tokens.clamp(min=1),
+        )  # n_ans
+        a2b = self.lm_predictions.lm_head.bias[a2tok]
+        self.answer_bias.weight = torch.div(
+            (a2b * (a2tok != pad_token_id).float()).sum(1), sum_tokens.clamp(min=1)
+        )
+        if freeze_last:
+            self.answer_embeddings.requires_grad_(False)
+            self.answer_bias.requires_grad_(False)
+
+    def emd_context_layer(self, encoder_layers, z_states, attention_mask, encoder, temporal_factor, train_mode):
+        if attention_mask.dim() <= 2:
+            extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            att_mask = extended_attention_mask.byte()
+            attention_mask = att_mask * att_mask.squeeze(-2).unsqueeze(-1)
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask.unsqueeze(1)
+        hidden_states = encoder_layers[-2]
+        if not self.config.position_biased_input:
+            layers = [encoder.layer[-1] for _ in range(2)]
+            z_states = z_states + hidden_states
+            query_states = z_states
+            query_mask = attention_mask
+            outputs = []
+            rel_embeddings = encoder.get_rel_embedding()
+
+            for layer in layers:
+                output = layer(
+                    temporal_factor,
+                    hidden_states,
+                    query_mask,
+                    return_att=False,
+                    query_states=query_states,
+                    relative_pos=None,
+                    rel_embeddings=rel_embeddings,
+                    train_mode=train_mode
+                )
+                query_states = output[0]
+                outputs.append(query_states)
+        else:
+            outputs = [encoder_layers[-1]]
+
+        return outputs
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        video=None,
+        video_mask=None,
+        train_mode=False,
+    ):
+        token_type_ids=None
+        position_ids=None
+        inputs_embeds=None
+        output_attentions=None
+        return_dict=None
+        mlm=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
+        )
+
+        
+        # rand_video = torch.randn(1,30,3,224,224).cuda()
+        # video = self.clip.encode_image(rand_video.squeeze()).unsqueeze(0)
+        # video = video.to(torch.float)
+
+        outputs = self.deberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=return_dict,
+            video=video,
+            video_mask=video_mask,
+            train_mode = train_mode
+        )
+
+        loss_moe = outputs['loss_moe']
+
+        if labels is not None:
+            if (
+                self.features_dim and video is not None
+            ):  # ignore the label predictions for visual tokens
+                video_shape = video[:, :, 0].size()
+                # add video_general
+                if self.add_video_feat:
+                    video_shape = (video_shape[0], video_shape[1] + 1)
+
+                video_labels = torch.tensor(
+                    [[-100] * video_shape[1]] * video_shape[0],
+                    dtype=torch.long,
+                    device=labels.device,
+                )
+                labels = torch.cat([video_labels, labels], 1)
+
+        # sequence_output = outputs[0]
+        modified = self.emd_context_layer(
+            encoder_layers=outputs["hidden_states"],
+            z_states=outputs["position_embeddings"].repeat(
+                input_ids.shape[0] // len(outputs["position_embeddings"]), 1, 1
+            ),
+            attention_mask=outputs["attention_mask"],
+            encoder=self.deberta.encoder,
+            temporal_factor=outputs["video_g"],
+            train_mode = train_mode
+        )
+        bias = None
+        if self.n_ans and (not mlm):  # downstream mode
+            embeddings = self.answer_embeddings.weight
+            bias = self.answer_bias
+        else:
+            embeddings = self.deberta.embeddings.word_embeddings.weight
+        prediction_scores = self.lm_predictions(modified[-1], embeddings, bias)
+
+        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),  # labels[labels > 0].view(-1)
+            )
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[1:]
+            return (
+                ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+            )
+        
+        return MaskedLMOutput(
+            loss_moe=loss_moe,
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            loads=outputs.loads,
+            embeddings=outputs.video_g
+        )
+
+
+# copied from transformers.models.bert.BertPredictionHeadTransform with bert -> deberta
+class DebertaV2PredictionHeadTransform(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
+
+
+# copied from transformers.models.bert.BertLMPredictionHead with bert -> deberta
+class DebertaV2LMPredictionHead(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)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+
+        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`
+
+    def forward(self, hidden_states, embedding_weight, bias=None):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        if bias is not None:
+            logits = (
+                torch.matmul(hidden_states, embedding_weight.t().to(hidden_states))
+                + bias
+            )
+        else:
+            logits = (
+                torch.matmul(hidden_states, embedding_weight.t().to(hidden_states))
+                + self.bias
+            )
+        return logits
+
+
+# copied from transformers.models.bert.BertOnlyMLMHead with bert -> deberta
+class DebertaV2OnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        # self.predictions = DebertaV2LMPredictionHead(config)
+        self.lm_head = DebertaV2LMPredictionHead(config)
+
+    def forward(self, sequence_output, embedding_weight, bias=None):
+        prediction_scores = self.lm_head(sequence_output, embedding_weight, bias=bias)
+        return prediction_scores

model/evl.py

@@ -0,0 +1,345 @@
+
+from typing import Dict, Iterable, List, Tuple
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from collections import OrderedDict
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class Attention(nn.Module):
+    '''
+    A generalized attention module with more flexibility.
+    '''
+
+    def __init__(
+        self, q_in_dim: int, k_in_dim: int, v_in_dim: int,
+        qk_proj_dim: int, v_proj_dim: int, num_heads: int, out_dim: int,
+        return_all_features: bool = False, add_mask: bool = False, dropout: float = 0.0
+    ):
+        super().__init__()
+
+        self.q_proj = nn.Linear(q_in_dim, qk_proj_dim)
+        self.k_proj = nn.Linear(k_in_dim, qk_proj_dim)
+        self.v_proj = nn.Linear(v_in_dim, v_proj_dim)
+        self.out_proj = nn.Linear(v_proj_dim, out_dim)
+
+        self.num_heads = num_heads
+        self.return_all_features = return_all_features
+        assert qk_proj_dim % num_heads == 0 and v_proj_dim % num_heads == 0
+
+        self.add_mask = add_mask
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj):
+            nn.init.xavier_uniform_(m.weight)
+            nn.init.constant_(m.bias, 0.)
+
+    def forward(self, q, k, v, mask):
+        if not self.add_mask:
+            mask = torch.ones_like(mask)
+
+        assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3
+        N = q.size(0); assert k.size(0) == N and v.size(0) == N
+        Lq, Lkv = q.size(1), k.size(1); assert v.size(1) == Lkv
+
+        q, k, v = self.q_proj(q), self.k_proj(k), self.v_proj(v)
+        
+        H = self.num_heads
+        Cqk, Cv = q.size(-1) // H, v.size(-1) // H
+
+        q = q.view(N, Lq, H, Cqk)
+        k = k.view(N, Lkv, H, Cqk)
+        v = v.view(N, Lkv, H, Cv)
+
+        aff = torch.einsum('nqhc,nkhc->nqkh', q / (Cqk ** 0.5), k)
+        #aff = aff.softmax(dim=-2)
+        
+        rmask = ~(mask.bool())
+        aff = aff.masked_fill(rmask.unsqueeze(1).unsqueeze(-1).to(aff.device), float("-inf"))
+        aff = aff.softmax(dim = -2)
+
+        mix = torch.einsum('nqlh,nlhc->nqhc', aff, v)
+
+        out = self.out_proj(mix.flatten(-2))
+
+        if self.return_all_features:
+            return dict(q=q, k=k, v=v, aff=aff, out=out)
+        else:
+            return out
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        in_feature_dim: int = 768,
+        qkv_dim: int = 768,
+        num_heads: int = 12,
+        mlp_factor: float = 4.0,
+        mlp_dropout: float = 0.0,
+        act: nn.Module = QuickGELU,
+        add_mask: bool = False
+    ):
+        super().__init__()
+
+        self.attn = Attention(
+            q_in_dim=in_feature_dim, k_in_dim=in_feature_dim, v_in_dim=in_feature_dim,
+            qk_proj_dim=qkv_dim, v_proj_dim=qkv_dim, num_heads=num_heads, out_dim=in_feature_dim, add_mask=add_mask
+        )
+
+        mlp_dim = round(mlp_factor * in_feature_dim)
+        self.mlp = nn.Sequential(OrderedDict([
+            ('fc1', nn.Linear(in_feature_dim, mlp_dim)),
+            ('act', act()),
+            ('dropout', nn.Dropout(mlp_dropout)),
+            ('fc2', nn.Linear(mlp_dim, in_feature_dim)),
+        ]))
+
+        self.norm1 = LayerNorm(in_feature_dim)
+        self.norm2 = LayerNorm(in_feature_dim)
+        self.norm3 = LayerNorm(in_feature_dim)
+
+        self._initialize_weights()
+
+
+    def _initialize_weights(self):
+        for m in (self.mlp[0], self.mlp[-1]):
+            nn.init.xavier_uniform_(m.weight)
+            nn.init.normal_(m.bias, std=1e-6)
+
+
+    def forward(self, x, y, mask):
+        y_norm = self.norm3(y)
+        x = x + self.attn(self.norm1(x), y_norm, y_norm, mask)
+        x = x + self.mlp(self.norm2(x))
+
+        return x
+
+
+class EVLDecoder(nn.Module):
+
+    def __init__(
+        self,
+        num_frames: int = 8,
+        spatial_size: Tuple[int, int] = (14, 14),
+        num_layers: int = 4,
+        in_feature_dim: int = 768,
+        qkv_dim: int = 768,
+        num_heads: int = 12,
+        mlp_factor: float = 4.0,
+        enable_temporal_conv: bool = True,
+        enable_temporal_pos_embed: bool = True,
+        mlp_dropout: float = 0.5,
+        add_vid_feat: bool = False,
+        add_mask: bool = False,
+    ):
+        super().__init__()
+
+        self.num_layers = num_layers
+
+        self.add_vid_feat = add_vid_feat
+
+        if add_vid_feat:
+            self.decoder_layers = nn.ModuleList(
+                [TransformerDecoderLayer(in_feature_dim, qkv_dim, num_heads, mlp_factor, mlp_dropout, add_mask=add_mask) for _ in range(num_layers)]
+            )
+            self.cls_token = nn.Parameter(torch.zeros([in_feature_dim]))
+            self._initialize_weights()
+
+        if enable_temporal_conv:
+            self.temporal_conv = nn.ModuleList(
+                [nn.Conv1d(in_feature_dim, in_feature_dim, kernel_size=3, stride=1, padding=1, groups=in_feature_dim) for _ in range(num_layers)]
+            )
+
+            # self.temporal_conv = nn.ModuleList(
+            #     [nn.Linear(in_feature_dim, in_feature_dim) for _ in range(num_layers)]
+            # )
+
+        if enable_temporal_pos_embed:
+            self.temporal_pos_embed = nn.ParameterList(
+                [nn.Parameter(torch.zeros([num_frames, in_feature_dim])) for _ in range(num_layers)]
+            )
+
+    def _initialize_weights(self):
+        nn.init.normal_(self.cls_token, std=0.02)
+
+    def forward(self, in_features, video_mask):
+        N, T, C = in_features.size()
+        
+        if self.add_vid_feat:
+            x = self.cls_token.view(1, 1, -1).repeat(N, 1, 1)
+
+        frame_features = in_features
+        for i in range(self.num_layers):
+            frame_features = in_features
+            feat = in_features
+
+            feat = feat.permute(0, 2, 1).contiguous() # N * L, C, T
+            
+            
+            feat = self.temporal_conv[i](feat)
+            feat = feat.view(N, C, T).permute(0, 2, 1,).contiguous() # N, T, C
+            frame_features = frame_features + feat
+        
+            frame_features = frame_features + self.temporal_pos_embed[i].view(1, T, C)
+            
+            if self.add_vid_feat:
+                x = self.decoder_layers[i](x, frame_features, video_mask)
+        
+        if self.add_vid_feat:
+            return x
+        
+        return frame_features
+
+
+class EVLTransformer(nn.Module):
+
+    def __init__(
+        self,
+        num_frames: int = 8,
+        decoder_num_layers: int = 2,
+        decoder_qkv_dim: int = 768,
+        decoder_num_heads: int = 16,
+        decoder_mlp_factor: float = 4.0,
+        enable_temporal_conv: bool = True,
+        enable_temporal_pos_embed: bool = True,
+        enable_temporal_cross_attention: bool = False,
+        decoder_mlp_dropout: float = 0.5,
+        add_video_feat: bool = False,
+        output_dim: int = 1536,
+        add_mask: bool = False
+    ):
+        super().__init__()
+
+        self.decoder_num_layers = decoder_num_layers
+
+        backbone_feature_dim = 768
+        backbone_spatial_size = (16, 16)
+
+        self.decoder = EVLDecoder(
+            num_frames=num_frames,
+            spatial_size=backbone_spatial_size,
+            num_layers=decoder_num_layers,
+            in_feature_dim=backbone_feature_dim,
+            qkv_dim=decoder_qkv_dim,
+            num_heads=decoder_num_heads,
+            mlp_factor=decoder_mlp_factor,
+            enable_temporal_conv=enable_temporal_conv,
+            enable_temporal_pos_embed=enable_temporal_pos_embed,
+            mlp_dropout=decoder_mlp_dropout,
+            add_vid_feat = add_video_feat,
+            add_mask=add_mask
+        )
+        self.add_vid_feat = add_video_feat
+        if self.add_vid_feat:
+            self.norm = nn.LayerNorm(backbone_feature_dim)
+            #self.dropout = nn.Dropout(0.5)
+            self.proj = nn.Linear(decoder_qkv_dim, output_dim)
+
+    def forward(self, x, video_mask):
+
+        features = x
+        x = self.decoder(features, video_mask)
+        if self.add_vid_feat:
+            x = self.norm(x)
+            #x = self.dropout(x)
+            x = self.proj(x)
+
+        return x
+
+class TemporalAttention(nn.Module):
+    def __init__(
+        self,
+        in_feature_dim: int = 768,
+        qkv_dim: int = 768,
+        num_heads: int = 8,
+        max_frames: int = 40,
+        stride: int = 4,
+        kernel_size: int = 4,
+        add_mask: bool = True,
+    ):
+        super().__init__()
+        
+        self.num_layers = 2
+        self.kernel_size = kernel_size
+        self.stride = stride
+        max_frames = (max_frames - self.kernel_size) // self.stride + 1
+        
+        self.decoder_layers = nn.ModuleList(
+                [TransformerDecoderLayer(in_feature_dim, qkv_dim, num_heads, 2.0, 0.5, add_mask=add_mask) for _ in range(self.num_layers)]
+            )
+        
+        '''
+        self.attn = Attention(
+            q_in_dim=in_feature_dim, k_in_dim=in_feature_dim, v_in_dim=in_feature_dim,
+            qk_proj_dim=qkv_dim, v_proj_dim=qkv_dim, num_heads=num_heads, out_dim=in_feature_dim, add_mask=add_mask
+        )'''
+
+        self.temporal_pos_embed = nn.Parameter(torch.zeros([max_frames, in_feature_dim]))
+        self.norm = nn.LayerNorm(in_feature_dim)
+    
+    def forward(self, x, video_mask):
+        
+        x, video_mask = avg_1d_pool(x, self.kernel_size, self.stride, video_mask, return_mask=True)
+
+        x = x + self.temporal_pos_embed.unsqueeze(0)
+        for i in range(self.num_layers):
+            x = self.decoder_layers[i](x, x, video_mask)
+
+        #x_norm = self.norm(x)
+        #x = x + self.attn(x_norm, x_norm, x_norm, video_mask)
+        
+        return x
+    
+def recursive_gumbel_softmax(sim, x, video_mask, topk):
+    # sim: bs, T
+    # x: bs, T, dim
+
+    feats = []
+    bs = x.shape[0]
+    idxs = torch.zeros(bs, 10)
+    v_masks = []
+
+    rmask = ~(video_mask.bool())
+    sim = sim.masked_fill(rmask.unsqueeze(1).to(sim.device), float("-inf"))
+
+    for i in range(topk):
+        choice = F.gumbel_softmax(sim/0.01, hard=True, dim = -1, tau=0.1).squeeze(1) # bs, T
+        idxs[:, i] = torch.argsort(choice, descending=True)[:, 0]
+        tmp = torch.sum(choice.unsqueeze(-1) * x, dim = 1, keepdim=True) # bs, dim
+        feats.append(tmp)
+
+        mask_tmp = video_mask[torch.arange(bs), idxs[:, i].to(torch.long)]
+        v_masks.append(mask_tmp)
+        sim = sim - choice.unsqueeze(1)
+    
+    rank = torch.argsort(idxs, dim = 1)
+    
+    feats = torch.cat(feats, dim= 1) # bs, 10, dim
+    res = [feats[torch.arange(bs), rank[:, i]] for i in range(10)]
+    res = torch.stack(res, dim=1)
+    
+    video_mask = torch.stack(v_masks, dim=1)
+    video_mask = [video_mask[torch.arange(bs), rank[:, i]] for i in range(10)]
+    video_mask = torch.stack(video_mask, dim = 1)
+
+    return res, video_mask
+ 
+

model/moe.py

@@ -0,0 +1,442 @@
+
+### copy from LIMoE
+
+
+#from distutils.command.config import config
+import os
+import torch
+import torch.nn as nn
+from torch.distributions.normal import Normal
+import torch.nn.functional as F
+import numpy as np
+
+from transformers.activations import ACT2FN
+from .adapter import Adapter
+from collections import OrderedDict
+from copy import deepcopy
+
+#-------------------#
+# MoE
+
+class MLP(nn.Module):
+    def __init__(self, input_size:int, output_size:int, hidden_size:int):
+        super(MLP, self).__init__()
+        self.fc1 = nn.Linear(input_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, output_size)
+        self.dropout = nn.Dropout(0.1)
+        self.activation = ACT2FN["gelu"]
+        self.log_soft = nn.LogSoftmax(1)
+        self.apply(self.init_weights)
+
+    def init_weights(self, m: nn.Module, std=1e-3):
+        if isinstance(m, nn.Linear):
+            torch.nn.init.normal_(m.weight, std=std)
+            torch.nn.init.normal_(m.bias, std=std)
+            m.weight.data = torch.clamp(m.weight.data, min=-2 * std, max=2 * std)
+            m.bias.data = torch.clamp(m.bias.data, min=-2 * std, max=2 * std)
+        elif isinstance(m, nn.LayerNorm):
+            m.bias.data.zero_()
+            m.weight.data.fill_(1.0)
+
+    def forward(self, x):
+        out = self.fc1(x)
+        out = self.activation(out)
+        out = self.dropout(out)
+        out = self.fc2(out)
+        out = self.log_soft(out)
+        return out
+
+
+
+class SparseDispatcher(object):
+    """Helper for implementing a mixture of experts.
+    The purpose of this class is to create input minibatches for the
+    experts and to combine the results of the experts to form a unified
+    output tensor.
+    There are two functions:
+    dispatch - take an input Tensor and create input Tensors for each expert.
+    combine - take output Tensors from each expert and form a combined output
+      Tensor.  Outputs from different experts for the same batch element are
+      summed together, weighted by the provided "gates".
+    The class is initialized with a "gates" Tensor, which specifies which
+    batch elements go to which experts, and the weights to use when combining
+    the outputs.  Batch element b is sent to expert e iff gates[b, e] != 0.
+    The inputs and outputs are all two-dimensional [batch, depth].
+    Caller is responsible for collapsing additional dimensions prior to
+    calling this class and reshaping the output to the original shape.
+    See common_layers.reshape_like().
+    Example use:
+    gates: a float32 `Tensor` with shape `[batch_size, num_experts]`
+    inputs: a float32 `Tensor` with shape `[batch_size, input_size]`
+    experts: a list of length `num_experts` containing sub-networks.
+    dispatcher = SparseDispatcher(num_experts, gates)
+    expert_inputs = dispatcher.dispatch(inputs)
+    expert_outputs = [experts[i](expert_inputs[i]) for i in range(num_experts)]
+    outputs = dispatcher.combine(expert_outputs)
+    The preceding code sets the output for a particular example b to:
+    output[b] = Sum_i(gates[b, i] * experts[i](inputs[b]))
+    This class takes advantage of sparsity in the gate matrix by including in the
+    `Tensor`s for expert i only the batch elements for which `gates[b, i] > 0`.
+    """
+
+    def __init__(self, num_experts, gates):
+        """Create a SparseDispatcher."""
+
+        self._gates = gates
+        self._num_experts = num_experts
+        # sort experts
+        sorted_experts, index_sorted_experts = torch.nonzero(gates).sort(0)  # torch.nonzero: 返回非0坐标，按行、列依次排序
+        # drop indices
+        _, self._expert_index = sorted_experts.split(1, dim=1)
+        # get according batch index for each expert
+        self._batch_index = sorted_experts[index_sorted_experts[:, 1],0]
+        # calculate num samples that each expert gets
+        self._part_sizes = list((gates > 0).sum(0).cpu().numpy())
+        # expand gates to match with self._batch_index
+        gates_exp = gates[self._batch_index.flatten()]
+        self._nonzero_gates = torch.gather(gates_exp, 1, self._expert_index)
+
+    def dispatch(self, inp):
+        """Create one input Tensor for each expert.
+        The `Tensor` for a expert `i` contains the slices of `inp` corresponding
+        to the batch elements `b` where `gates[b, i] > 0`.
+        Args:
+          inp: a `Tensor` of shape "[batch_size, <extra_input_dims>]`
+        Returns:
+          a list of `num_experts` `Tensor`s with shapes
+            `[expert_batch_size_i, <extra_input_dims>]`.
+        """
+
+        # assigns samples to experts whose gate is nonzero
+
+        # expand according to batch index so we can just split by _part_sizes
+        inp_exp = inp[self._batch_index].squeeze(1)
+        return torch.split(inp_exp, self._part_sizes, dim=0)
+
+
+    def combine(self, expert_out, multiply_by_gates=True):
+        """Sum together the expert output, weighted by the gates.
+        The slice corresponding to a particular batch element `b` is computed
+        as the sum over all experts `i` of the expert output, weighted by the
+        corresponding gate values.  If `multiply_by_gates` is set to False, the
+        gate values are ignored.
+        Args:
+          expert_out: a list of `num_experts` `Tensor`s, each with shape
+            `[expert_batch_size_i, <extra_output_dims>]`.
+          multiply_by_gates: a boolean
+        Returns:
+          a `Tensor` with shape `[batch_size, <extra_output_dims>]`.
+        """
+        # apply exp to expert outputs, so we are not longer in log space
+        
+        #stitched = torch.cat(expert_out, 0).exp()
+        stitched = torch.cat(expert_out, 0)
+
+        if multiply_by_gates:
+            if len(stitched.shape) == 3:
+                stitched = stitched.mul(self._nonzero_gates.unsqueeze(1))
+            else:
+                stitched = stitched.mul(self._nonzero_gates)
+
+        if len(stitched.shape) == 3:
+            zeros = torch.zeros(self._gates.size(0), expert_out[-1].size(1), expert_out[-1].size(-1), requires_grad=True, device=stitched.device)
+        else:
+            zeros = torch.zeros(self._gates.size(0), expert_out[-1].size(1), requires_grad=True, device=stitched.device)
+        # combine samples that have been processed by the same k experts
+        combined = zeros.index_add(0, self._batch_index, stitched.float())
+        # add eps to all zero values in order to avoid nans when going back to log space
+        
+        #combined[combined == 0] = np.finfo(float).eps
+        # back to log space
+        #return combined.log()
+        return combined
+
+
+    def expert_to_gates(self):
+        """Gate values corresponding to the examples in the per-expert `Tensor`s.
+        Returns:
+          a list of `num_experts` one-dimensional `Tensor`s with type `tf.float32`
+              and shapes `[expert_batch_size_i]`
+        """
+        # split nonzero gates for each expert
+        return torch.split(self._nonzero_gates, self._part_sizes, dim=0)
+
+
+
+
+class MoE(nn.Module):
+
+    """Call a Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
+    Args:
+    input_size: integer - size of the input
+    output_size: integer - size of the input
+    num_experts: an integer - number of experts
+    hidden_size: an integer - hidden size of the experts
+    noisy_gating: a boolean
+    k: an integer - how many experts to use for each batch element
+    """
+
+    def __init__(self, 
+                 noisy_gating = True, 
+                 ds_factor = 8.0,
+                 num_experts = 4,
+                 moe_input_size = 768,
+                 top_k = 2,
+                 dropout = 0.1,
+                 gating = 'linear',
+                 routing = None,
+                 layer_id = 0
+                 ):
+        super(MoE, self).__init__()
+        self.noisy_gating = noisy_gating
+        self.num_experts = num_experts
+        self.input_size = moe_input_size
+        self.k = top_k
+        self.layer_id = layer_id
+
+        # instantiate experts
+        #self.experts = nn.ModuleList([MLP(self.input_size, self.output_size, self.hidden_size) for i in range(self.num_experts)])
+        self.gating = gating
+        self.experts = nn.ModuleList([Adapter(ds_factor, moe_input_size, dropout=dropout) for i in range(self.num_experts)])
+        self.routing = routing
+        self.infer_expert = None
+
+
+        if self.routing != 'random':
+            if gating == 'linear':
+                #self.w_gate = nn.Linear(self.input_size, self.num_experts, bias=False)
+                self.w_gate = nn.Parameter(torch.zeros(self.input_size, num_experts), requires_grad=True)
+            elif gating == 'cosine':
+                self.w_gate = CosineTopKGate(self.input_size, self.num_experts)
+            self.w_noise = nn.Parameter(torch.zeros(self.input_size, self.num_experts), requires_grad=True)
+
+            self.softplus = nn.Softplus()
+            self.softmax = nn.Softmax(-1)
+            self.register_buffer("mean", torch.tensor([0.0]))
+            self.register_buffer("std", torch.tensor([1.0]))
+
+        assert(self.k <= self.num_experts)
+
+    def cv_squared(self, x):
+        """The squared coefficient of variation of a sample.
+        Useful as a loss to encourage a positive distribution to be more uniform.
+        Epsilons added for numerical stability.
+        Returns 0 for an empty Tensor.
+        Args:
+        x: a `Tensor`.
+        Returns:
+        a `Scalar`.
+        """
+        eps = 1e-10
+        # if only num_experts = 1
+        if x.shape[0] == 1:
+            return torch.Tensor([0])
+        if len(x.shape) == 2:
+            x = x.sum(dim=0)
+        return x.float().var() / (x.float().mean()**2 + eps)
+
+
+    def _gates_to_load(self, gates):
+        """Compute the true load per expert, given the gates.
+        The load is the number of examples for which the corresponding gate is >0.
+        Args:
+        gates: a `Tensor` of shape [batch_size, n]
+        Returns:
+        a float32 `Tensor` of shape [n]
+        """
+        return (gates > 0).sum(0)
+
+
+    def _prob_in_top_k(self, clean_values, noisy_values, noise_stddev, noisy_top_values):
+        """Helper function to NoisyTopKGating.
+        Computes the probability that value is in top k, given different random noise.
+        This gives us a way of backpropagating from a loss that balances the number
+        of times each expert is in the top k experts per example.
+        In the case of no noise, pass in None for noise_stddev, and the result will
+        not be differentiable.
+        Args:
+        clean_values: a `Tensor` of shape [batch, n].
+        noisy_values: a `Tensor` of shape [batch, n].  Equal to clean values plus
+          normally distributed noise with standard deviation noise_stddev.
+        noise_stddev: a `Tensor` of shape [batch, n], or None
+        noisy_top_values: a `Tensor` of shape [batch, m].
+           "values" Output of tf.top_k(noisy_top_values, m).  m >= k+1
+        Returns:
+        a `Tensor` of shape [batch, n].
+        """
+
+        batch = clean_values.size(0)
+        m = noisy_top_values.size(1)
+        top_values_flat = noisy_top_values.flatten() # (bs x m)
+        threshold_positions_if_in = torch.arange(batch) * m + self.k # bs
+        threshold_if_in = torch.unsqueeze(torch.gather(top_values_flat, 0, threshold_positions_if_in.to(top_values_flat.device)), 1)
+
+        if len(noisy_values.shape) == 3:
+            threshold_if_in = threshold_if_in.unsqueeze(1)
+
+        is_in = torch.gt(noisy_values, threshold_if_in)
+        threshold_positions_if_out = threshold_positions_if_in - 1
+        threshold_if_out = torch.unsqueeze(torch.gather(top_values_flat,0 , threshold_positions_if_out.to(top_values_flat.device)), 1)
+        if len(noisy_values.shape) == 3:
+            threshold_if_out = threshold_if_out.unsqueeze(1)
+
+        # is each value currently in the top k.
+
+        normal = Normal(self.mean.to(noise_stddev.device), self.std.to(noise_stddev.device))
+        prob_if_in = normal.cdf((clean_values - threshold_if_in)/noise_stddev)
+        prob_if_out = normal.cdf((clean_values - threshold_if_out)/noise_stddev)
+        prob = torch.where(is_in, prob_if_in, prob_if_out)
+        return prob
+
+    
+    def random_k_gating(self, features, train):
+        if train:
+            idx = torch.randint(0, self.num_experts, 1)
+            results = self.experts[idx](features)
+
+        else:
+            results = []
+            for i in range(self.num_experts):
+                tmp = self.num_experts[i](features)
+                results.append(tmp)
+            
+            results = torch.stack(results, dim=0).mean(dim=0)
+
+        return results
+
+
+
+    def noisy_top_k_gating(self, x, train, noise_epsilon=1e-2):
+        """Noisy top-k gating.
+          See paper: https://arxiv.org/abs/1701.06538.
+          Args:
+            x: input Tensor with shape [batch_size, input_size]
+            train: a boolean - we only add noise at training time.
+            noise_epsilon: a float
+          Returns:
+            gates: a Tensor with shape [batch_size, num_experts]
+            load: a Tensor with shape [num_experts]
+        """
+        #clean_logits = self.w_gate(x)
+        if self.gating == 'linear':
+            clean_logits = x @ self.w_gate
+        elif self.gating == 'cosine':
+            clean_logits = self.w_gate(x)
+
+        if self.noisy_gating and train:
+            raw_noise_stddev = x @ self.w_noise
+            noise_stddev = ((self.softplus(raw_noise_stddev) + noise_epsilon) * train)
+            noisy_logits = clean_logits + ( torch.randn_like(clean_logits) * noise_stddev)
+            logits = noisy_logits
+        else:
+            logits = clean_logits
+
+        # logits (bs, n): 表示选择n中每个expert的概率
+
+        # 选k个experts，返回相应的下标以及logit
+        top_logits, top_indices = logits.topk(min(self.k + 1, self.num_experts), dim= -1)  
+        
+        top_k_logits = top_logits[:, :self.k] if len(top_logits.shape) == 2 else top_logits[:, :, :self.k]    
+        top_k_indices = top_indices[:, :self.k] if len(top_indices.shape) == 2 else top_indices[:, :, :self.k]
+        
+        top_k_gates = self.softmax(top_k_logits)
+
+        zeros = torch.zeros_like(logits, requires_grad=True)
+        # 将经过softmax后的weight分配给相应的expert，未选定的expert的weight则为0
+        gates = zeros.scatter(-1, top_k_indices, top_k_gates)  
+
+        if self.noisy_gating and self.k < self.num_experts and train:
+            load = (self._prob_in_top_k(clean_logits, noisy_logits, noise_stddev, top_logits)).sum(0)
+        else:
+            load = self._gates_to_load(gates)
+        return gates, load
+
+
+    def forward(self, x, frame_features, train=True, loss_coef=1e-2):
+        """Args:
+        x: tensor shape [batch_size, input_size]
+        train: a boolean scalar.
+        loss_coef: a scalar - multiplier on load-balancing losses
+        Returns:
+        y: a tensor with shape [batch_size, output_size].
+        extra_training_loss: a scalar.  This should be added into the overall
+        training loss of the model.  The backpropagation of this loss
+        encourages all experts to be approximately equally used across a batch.
+        """
+
+        if self.routing == 'random':
+            loss = None
+            load = None
+            if train:
+                gates = torch.zeros(x.shape[0], self.num_experts)
+                random_idx = torch.randint(0, self.num_experts, (x.shape[0],))
+                gates[torch.arange(x.shape[0]), random_idx] = 1
+                gates = gates.to(x.device)
+                dispatcher = SparseDispatcher(self.num_experts, gates)
+            
+                expert_inputs = dispatcher.dispatch(frame_features)  # 获取每个expert的输入
+                gates = dispatcher.expert_to_gates() #
+                expert_outputs = [self.experts[i](expert_inputs[i]) for i in range(self.num_experts)]
+                y = dispatcher.combine(expert_outputs)
+            else:
+                if self.infer_expert is None:
+                    weights = [self.experts[i].state_dict() for i in range(self.num_experts)]
+                    merge_weights = OrderedDict()
+                    for idx, it in enumerate(weights):
+                        for k,v in it.items():
+                            merge_weights[k] = v / self.num_experts if idx==0 else merge_weights[k] + v / self.num_experts
+                
+                    self.infer_expert = deepcopy(self.experts[0])
+                    self.infer_expert.load_state_dict(merge_weights)
+                
+                y = self.infer_expert(frame_features)
+            
+            return y, loss, load 
+        
+        else:
+            if len(x.shape) == 1:
+                x = x.unsqueeze(0)
+            
+            gates, load = self.noisy_top_k_gating(x, train)
+
+            # calculate importance loss
+            importance = gates.sum(dim=0)
+            
+            # calculate loss
+            loss = self.cv_squared(importance) + self.cv_squared(load)
+            loss *= loss_coef
+
+            dispatcher = SparseDispatcher(self.num_experts, gates)
+            
+            expert_inputs = dispatcher.dispatch(frame_features)  # 获取每个expert的输入
+            gates = dispatcher.expert_to_gates() # 获取
+            expert_outputs = [self.experts[i](expert_inputs[i]) for i in range(self.num_experts)]
+            y = dispatcher.combine(expert_outputs)
+            return y, loss, load
+
+class CosineTopKGate(torch.nn.Module):
+    def __init__(self, model_dim, num_global_experts, proj_dim=256, init_t=0.5):
+        super(CosineTopKGate, self).__init__()
+        self.temperature = torch.nn.Parameter(torch.log(torch.full([1], 1.0 / init_t)), requires_grad=True)
+        self.cosine_projector = torch.nn.Linear(model_dim, proj_dim)
+        self.sim_matrix = torch.nn.Parameter(torch.randn(size=(proj_dim, num_global_experts)), requires_grad=True)
+        self.clamp_max = torch.log(torch.tensor(1. / 0.01)).item()
+        torch.nn.init.normal_(self.sim_matrix, 0, 0.01)
+
+    def forward(self, x):
+        cosine_projector = self.cosine_projector
+        sim_matrix = self.sim_matrix
+        logits = torch.matmul(F.normalize(cosine_projector(x), dim=1),
+                              F.normalize(sim_matrix, dim=0))
+        logit_scale = torch.clamp(self.temperature, max=self.clamp_max).exp()
+        logits = logits * logit_scale
+        return logits
+
+'''
+model = MoE()
+
+inputs = torch.randn((32, 1, 768))
+frame_features = torch.randn((32,10, 768))
+
+model(inputs, frame_features)
+'''

tmp.py

@@ -0,0 +1,30 @@
+import pandas as pd
+import json
+import glob
+
+result = json.load(open("/home/qinyixin/workspace/TgMoE/Frozenbilm/results/T_MoENet_NEXT-QA.json"))
+video_dir = "/mnt/hdd3/qinyixin/nextqa/video"
+
+cols = pd.read_csv("/mnt/hdd3/qinyixin/FrozenBilm/NEXT-QA/val.csv").columns.to_list()
+nextqa = pd.read_csv("/mnt/hdd3/qinyixin/FrozenBilm/NEXT-QA/val.csv").values
+qid_to_vidid = {}
+for it in nextqa:
+    choices = [it[9 + idx] for idx in range(5)]
+    answer = choices[it[6]]
+    question = it[5]
+    qid = it[7]
+    vidid = str(it[1])
+    vid_path = glob.glob(video_dir + "/*/"+ vidid + ".mp4")
+
+    qid_to_vidid[str(qid)] = {"vid_path": vid_path, 
+                              "choices": str(choices), 
+                              "question": question,
+                              "answer": answer
+                              }
+
+correct = []
+for k, v in result.items():
+    if v['acc']:
+        correct.append(qid_to_vidid[k])
+
+json.dump(correct, open("demo/T-MoENet_result.json", "w"))

tmp2.py

@@ -0,0 +1,10 @@
+import torch
+from Infer import Infer
+
+device = "cuda"
+handler = Infer(device)
+candidates = ['adjust the tree', 'get away the dust', 'dancing', 'pressed a button to activate', 'presents']
+with torch.no_grad():
+    handler.generate("why did the boy clap his hands when he ran to the christmas tree?", 
+                    "/home/qinyixin/workspace/TgMoE/Frozenbilm/demo/videos/4882821564.mp4",
+                    candidates)

videos/3249402410.mp4

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

videos/4882821564.mp4

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

videos/6233408665.mp4

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