crfm-helm 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

helm/clients/image_generation/aleph_alpha_image_generation_client.py

@@ -0,0 +1,98 @@
+from typing import List, Dict
+
+from helm.common.cache import Cache, CacheConfig
+from helm.common.request import Request, RequestResult, GeneratedOutput
+from helm.common.tokenization_request import (
+    TokenizationRequest,
+    TokenizationRequestResult,
+    DecodeRequest,
+    DecodeRequestResult,
+)
+from helm.clients.client import Client, CachingClient
+from .image_generation_client_utils import get_single_image_multimedia_object
+
+
+class AlephAlphaImageGenerationClient(Client):
+    """
+    Client for Aleph Alpha vision models. Offline eval only.
+    """
+
+    DEFAULT_IMAGE_HEIGHT: int = 512
+    DEFAULT_IMAGE_WIDTH: int = 512
+
+    DEFAULT_GUIDANCE_SCALE: float = 7.5
+    DEFAULT_STEPS: int = 50
+
+    @staticmethod
+    def convert_to_raw_request(request: Request) -> Dict:
+        raw_request: Dict = {
+            "request_type": "image-model-inference",
+            "model": request.model_engine,
+            "prompt": request.prompt,
+            "n": request.num_completions,
+            "guidance_scale": AlephAlphaImageGenerationClient.DEFAULT_GUIDANCE_SCALE,
+            "steps": AlephAlphaImageGenerationClient.DEFAULT_STEPS,
+            "width": AlephAlphaImageGenerationClient.DEFAULT_IMAGE_WIDTH,
+            "height": AlephAlphaImageGenerationClient.DEFAULT_IMAGE_HEIGHT,
+        }
+        if request.random is not None:
+            raw_request["random"] = request.random
+
+        assert request.image_generation_parameters is not None
+        if request.image_generation_parameters.guidance_scale is not None:
+            raw_request["guidance_scale"] = request.image_generation_parameters.guidance_scale
+        if request.image_generation_parameters.diffusion_denoising_steps is not None:
+            raw_request["steps"] = request.image_generation_parameters.diffusion_denoising_steps
+        if (
+            request.image_generation_parameters.output_image_width is not None
+            and request.image_generation_parameters.output_image_height is not None
+        ):
+            raw_request["width"] = request.image_generation_parameters.output_image_width
+            raw_request["height"] = request.image_generation_parameters.output_image_height
+
+        return raw_request
+
+    def __init__(self, cache_config: CacheConfig):
+        self._cache = Cache(cache_config)
+        self._promptist_model = None
+        self._promptist_tokenizer = None
+
+    def make_request(self, request: Request) -> RequestResult:
+        if request.model_engine != "m-vader":
+            raise ValueError(f"Unsupported model: {request.model_engine}")
+
+        raw_request = AlephAlphaImageGenerationClient.convert_to_raw_request(request)
+        raw_request.pop("random", None)
+        cache_key = CachingClient.make_cache_key(raw_request, request)
+
+        try:
+
+            def fail():
+                raise RuntimeError(
+                    f"The result has not been uploaded to the cache for the following request: {cache_key}"
+                )
+
+            response, cached = self._cache.get(cache_key, fail)
+        except RuntimeError as e:
+            error: str = f"AlephAlphaVisionClient error: {e}"
+            return RequestResult(success=False, cached=False, error=error, completions=[], embedding=[])
+
+        completions: List[GeneratedOutput] = [
+            GeneratedOutput(
+                text="", logprob=0, tokens=[], multimodal_content=get_single_image_multimedia_object(file_path)
+            )
+            for file_path in response["images"]
+        ]
+        return RequestResult(
+            success=True,
+            cached=cached,
+            request_time=response["request_time"],
+            completions=completions,
+            embedding=[],
+        )
+
+    def tokenize(self, request: TokenizationRequest) -> TokenizationRequestResult:
+        raise NotImplementedError("This client does not support tokenizing.")
+
+    def decode(self, request: DecodeRequest) -> DecodeRequestResult:
+        raise NotImplementedError("This client does not support decoding.")

helm/clients/image_generation/cogview2/coglm_strategy.py

@@ -0,0 +1,96 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   coglm_strategy.py
+@Time    :   2021/10/08 22:22:42
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import os
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+
+def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-65504):
+    # This function has been mostly taken from huggingface conversational ai code at
+    # https://medium.com/huggingface/how-to-build-a-state-of-the-art-conversational-ai-with-transfer-learning-2d818ac26313
+
+    if top_k > 0:
+        # Remove all tokens with a probability less than the last token of the top-k
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        # convert to 1D
+        logits = logits.view(logits.size()[1]).contiguous()
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+        # Remove tokens with cumulative probability above the threshold
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # Shift the indices to the right to keep also the first token above the threshold
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+        # going back to 2D
+        logits = logits.view(1, -1).contiguous()
+
+    return logits
+
+
+class CoglmStrategy:
+    def __init__(
+        self, invalid_slices=[], temperature=1.0, top_k=200, eps=1e-4, top_p=0.0, end_tokens=None, top_k_cluster=1.0
+    ):
+        self.invalid_slices = invalid_slices
+        self.temperature = temperature
+        self.topk = top_k
+        self.top_p = top_p
+        self.eps = eps
+        if end_tokens is None:
+            end_tokens = []
+        self.end_tokens = end_tokens
+        self._is_done = False
+        self.outlier_count_down = 5
+        self.cluster_labels = torch.tensor(
+            np.load(f"{os.path.dirname(os.path.abspath(__file__))}/cluster_label.npy"),
+            device="cuda" if torch.cuda.is_available() else "cpu",
+            dtype=torch.long,
+        )
+        self.top_k_cluster = top_k_cluster
+
+    @property
+    def is_done(self) -> bool:
+        return self._is_done
+
+    def forward(self, logits, tokens, mems, temperature=None):
+        if temperature is None:
+            temperature = self.temperature
+        logits = logits / temperature
+        for invalid_slice in self.invalid_slices:
+            logits[..., invalid_slice] = -65504
+
+        rprobs = F.softmax(logits.float(), dim=-1)
+        c = self.cluster_labels.expand(*rprobs.shape)
+        cprobs = torch.zeros(logits.shape[0], 500, device=logits.device).scatter_add_(1, c, rprobs)
+        best_scores, best_clusters = cprobs.topk(self.topk)
+        bz = logits.shape[0]
+        for i in range(bz):
+            best_scores[i] = best_scores[i]  # ** 0.2
+            selected_cluster = best_clusters[i][torch.multinomial(best_scores[i] / best_scores[i].sum(), num_samples=1)]
+            logits[i, self.cluster_labels != selected_cluster] = -65504
+
+        probs = F.softmax(logits.float() / self.top_k_cluster, dim=-1)  # float is essential, due to a bug in Pytorch
+        pred = torch.multinomial(probs, num_samples=1)
+
+        if pred.numel() == 1 and pred.item() in self.end_tokens:
+            self._is_done = True
+        tokens = torch.cat((tokens, pred.view(tokens.shape[0], 1)), dim=1)
+        return tokens, mems
+
+    def finalize(self, tokens, mems):
+        self._is_done = False
+        return tokens, mems

helm/clients/image_generation/cogview2/coglm_utils.py

@@ -0,0 +1,82 @@
+import torch
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from SwissArmyTransformer.model import CachedAutoregressiveModel
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+def get_masks_and_position_ids_coglm(seq, context_length):
+    tokens = seq.unsqueeze(0)
+    attention_mask = torch.ones((1, len(seq), len(seq)), device=tokens.device)
+    attention_mask.tril_()
+    attention_mask[..., :context_length] = 1
+    attention_mask.unsqueeze_(1)
+    position_ids = torch.zeros(len(seq), device=tokens.device, dtype=torch.long)
+    torch.arange(0, context_length, out=position_ids[:context_length])
+    torch.arange(512, 512 + len(seq) - context_length, out=position_ids[context_length:])
+    position_ids = position_ids.unsqueeze(0)
+    return tokens, attention_mask, position_ids
+
+
+def get_recipe(name):
+    r = {
+        "attn_plus": 1.4,
+        "temp_all_gen": 1.15,
+        "topk_gen": 16,
+        "temp_cluster_gen": 1.0,
+        "temp_all_dsr": 1.5,
+        "topk_dsr": 100,
+        "temp_cluster_dsr": 0.89,
+        "temp_all_itersr": 1.3,
+        "topk_itersr": 16,
+        "query_template": "{}<start_of_image>",
+    }
+    if name == "none":
+        pass
+    elif name == "mainbody":
+        r["query_template"] = "{} 高清摄影 隔绝<start_of_image>"
+    elif name == "photo":
+        r["query_template"] = "{} 高清摄影<start_of_image>"
+    elif name == "flat":
+        r["query_template"] = "{} 平面风格<start_of_image>"
+        # r['attn_plus'] = 1.8
+        # r['temp_cluster_gen'] = 0.75
+        r["temp_all_gen"] = 1.1
+        r["topk_dsr"] = 5
+        r["temp_cluster_dsr"] = 0.4
+        r["temp_all_itersr"] = 1
+        r["topk_itersr"] = 5
+    elif name == "comics":
+        r["query_template"] = "{} 漫画 隔绝<start_of_image>"
+        r["topk_dsr"] = 5
+        r["temp_cluster_dsr"] = 0.4
+        r["temp_all_gen"] = 1.1
+        r["temp_all_itersr"] = 1
+        r["topk_itersr"] = 5
+    elif name == "oil":
+        r["query_template"] = "{} 油画风格<start_of_image>"
+        pass
+    elif name == "sketch":
+        r["query_template"] = "{} 素描风格<start_of_image>"
+        r["temp_all_gen"] = 1.1
+    elif name == "isometric":
+        r["query_template"] = "{} 等距矢量图<start_of_image>"
+        r["temp_all_gen"] = 1.1
+    elif name == "chinese":
+        r["query_template"] = "{} 水墨国画<start_of_image>"
+        r["temp_all_gen"] = 1.12
+    elif name == "watercolor":
+        r["query_template"] = "{} 水彩画风格<start_of_image>"
+    return r
+
+
+class InferenceModel(CachedAutoregressiveModel):
+    def final_forward(self, logits, **kwargs):
+        logits_parallel = logits
+        logits_parallel = torch.nn.functional.linear(
+            logits_parallel.float(), self.transformer.word_embeddings.weight[:20000].float()
+        )
+        return logits_parallel

helm/clients/image_generation/cogview2/sr_pipeline/__init__.py

@@ -0,0 +1,15 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   __init__.py
+@Time    :   2022/03/02 13:57:09
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+from .direct_sr import DirectSuperResolution
+from .iterative_sr import IterativeSuperResolution
+from .sr_group import SRGroup
+
+DirectSuperResolution
+IterativeSuperResolution
+SRGroup

helm/clients/image_generation/cogview2/sr_pipeline/direct_sr.py

@@ -0,0 +1,96 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   inference_cogview2.py
+@Time    :   2021/10/10 16:31:34
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import torch
+from icetk import icetk as tokenizer
+
+from .dsr_sampling import filling_sequence_dsr, IterativeEntfilterStrategy
+from .dsr_model import DsrModel
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+
+class DirectSuperResolution:
+    def __init__(self, args, path, max_bz=4, shared_transformer=None):
+        try:
+            from SwissArmyTransformer.training.model_io import load_checkpoint
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        args.load = path
+        args.kernel_size = 5
+        args.kernel_size2 = 5
+        args.new_sequence_length = 4624
+        args.layout = [96, 496, 4096]
+
+        model = DsrModel(args, transformer=shared_transformer)
+        if args.fp16:
+            model = model.half()
+
+        load_checkpoint(model, args)  # on cpu
+        model.eval()
+        self.model = model.cuda() if torch.cuda.is_available() else model
+
+        # save cpu weights
+        self.saved_weights = dict((k, v.cpu()) for k, v in model.named_parameters() if "transformer" in k)
+
+        invalid_slices = [slice(tokenizer.num_image_tokens, None)]
+
+        self.strategy = IterativeEntfilterStrategy(
+            invalid_slices, temperature=args.temp_all_dsr, topk=args.topk_dsr, temperature2=args.temp_cluster_dsr
+        )  # temperature not used
+        self.max_bz = max_bz
+
+    def _restore_transformer_from_cpu(self, non_blocking=False):
+        for k, v in self.model.named_parameters():
+            if k in self.saved_weights:
+                v.copy_(self.saved_weights[k], non_blocking=non_blocking)
+
+    def __call__(self, text_tokens, image_tokens, enhance=False):
+        try:
+            from PIL import ImageEnhance, Image
+        except ModuleNotFoundError as e:
+            handle_module_not_found_error(e, ["heim"])
+
+        if len(text_tokens.shape) == 1:
+            text_tokens.unsqueeze_(0)
+        if len(image_tokens.shape) == 1:
+            image_tokens.unsqueeze_(0)
+
+        if enhance:
+            new_image_tokens = []
+            for small_img in image_tokens:
+                decoded = tokenizer.decode(image_ids=small_img).squeeze(0)
+                ndarr = decoded.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
+                image_pil_raw = ImageEnhance.Sharpness(Image.fromarray(ndarr))
+                small_img2 = tokenizer.encode(image_pil=image_pil_raw.enhance(1.0), image_size=160).view(-1)
+                new_image_tokens.append(small_img2)
+            image_tokens = torch.stack(new_image_tokens)
+
+        seq = torch.cat((text_tokens, image_tokens), dim=1)
+        seq1 = (
+            torch.tensor([tokenizer["<start_of_image>"]] * 3601, device=image_tokens.device)
+            .unsqueeze(0)
+            .expand(text_tokens.shape[0], -1)
+        )
+
+        self._restore_transformer_from_cpu()
+        model = self.model
+
+        output_list = []
+        for tim in range(max(text_tokens.shape[0] // self.max_bz, 1)):
+            output1 = filling_sequence_dsr(
+                model,
+                seq[tim * self.max_bz : (tim + 1) * self.max_bz],
+                seq1[tim * self.max_bz : (tim + 1) * self.max_bz],
+                warmup_steps=1,
+                block_hw=(1, 0),
+                strategy=self.strategy,
+            )
+            output_list.extend(output1[1:])
+        return torch.cat(output_list, dim=0)

helm/clients/image_generation/cogview2/sr_pipeline/dsr_model.py

@@ -0,0 +1,254 @@
+# -*- encoding: utf-8 -*-
+"""
+@File    :   cuda2d_model.py
+@Time    :   2021/10/02 01:36:32
+@Author  :   Ming Ding
+@Contact :   dm18@mails.tsinghua.edu.cn
+"""
+
+# here put the import lib
+import math
+import torch
+import torch.nn.functional as F
+
+from helm.common.optional_dependencies import handle_module_not_found_error
+
+try:
+    from SwissArmyTransformer.model.base_model import BaseModel, BaseMixin
+    from SwissArmyTransformer.mpu.utils import sqrt
+    from SwissArmyTransformer.mpu import ColumnParallelLinear, RowParallelLinear
+    from SwissArmyTransformer.model.transformer import unscaled_init_method, split_tensor_along_last_dim
+    from SwissArmyTransformer.ops.local_attention_function import f_similar, f_weighting
+except ModuleNotFoundError as e:
+    handle_module_not_found_error(e, ["heim"])
+
+
+class PositionEmbeddingMixin(BaseMixin):
+    def __init__(
+        self, additional_sequence_length, hidden_size, init_method_std=0.02, reinit_slice=slice(512, 512 + 400)
+    ):
+        super(PositionEmbeddingMixin, self).__init__()
+        self.reinit_slice = reinit_slice
+        self.position_embeddings = torch.nn.Embedding(additional_sequence_length, hidden_size)
+        torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std)
+
+    def reinit(self, parent_model=None):
+        old_weights = self.transformer.position_embeddings.weight.data[self.reinit_slice]
+        old_len, hidden_size = old_weights.shape
+        assert hidden_size == self.position_embeddings.weight.shape[-1]
+        old_edge, new_edge = sqrt(old_len), sqrt(self.position_embeddings.weight.shape[-2])
+        assert new_edge % old_edge == 0
+        self.position_embeddings.weight.data.view(
+            new_edge // old_edge, old_edge, new_edge // old_edge, old_edge, hidden_size
+        ).copy_(old_weights.view(1, old_edge, 1, old_edge, hidden_size))
+        # self.position_embeddings.weight.data.view(-1, old_len, hidden_size).copy_(old_weights)
+
+
+class AttentionMixin(BaseMixin):
+    def __init__(
+        self,
+        num_layers,
+        hidden_size,
+        init_method=unscaled_init_method(0.02),
+        output_layer_init_method=unscaled_init_method(0.02),
+    ):
+        super(AttentionMixin, self).__init__()
+        self.num_layers = num_layers  # replace attention in the LAST n layers
+        self.query_key_value = torch.nn.ModuleList(
+            [
+                ColumnParallelLinear(
+                    hidden_size, 3 * hidden_size, stride=3, gather_output=False, init_method=init_method
+                )
+                for layer_id in range(num_layers)
+            ]
+        )
+        self.dense = torch.nn.ModuleList(
+            [
+                RowParallelLinear(
+                    hidden_size, hidden_size, input_is_parallel=True, init_method=output_layer_init_method
+                )
+                for layer_id in range(num_layers)
+            ]
+        )
+
+    def reinit(self, parent_model=None):
+        start_layer = len(self.transformer.layers) - self.num_layers
+        assert start_layer >= 0
+        for layer_id in range(self.num_layers):
+            old_attention = self.transformer.layers[start_layer + layer_id].attention
+            self.query_key_value[layer_id].weight.data.copy_(old_attention.query_key_value.weight.data)
+            self.query_key_value[layer_id].bias.data.copy_(old_attention.query_key_value.bias.data)
+            self.dense[layer_id].weight.data.copy_(old_attention.dense.weight.data)
+            self.dense[layer_id].bias.data.copy_(old_attention.dense.bias.data)
+
+
+class DsrModel(BaseModel):
+    def __init__(self, args, transformer=None):
+        super().__init__(args, transformer=transformer)
+        self.original_sequence_length = args.max_sequence_length
+        additional_seqlen = args.new_sequence_length - args.max_sequence_length
+        self.add_mixin("extra_position_embedding", PositionEmbeddingMixin(additional_seqlen, args.hidden_size))
+        self.add_mixin("attention_plus", AttentionMixin(num_layers=args.num_layers, hidden_size=args.hidden_size))
+        self.layout = args.layout
+        # [PAD]... [ROI1] text ... [BOI1] {layout[0]} 1024 {layout[1]} [EOI1] 4095 {layout[2]}
+        self.kernel_size = args.kernel_size
+        self.kernel_size2 = args.kernel_size2
+        self.log_attention_weights = None
+
+    def position_embedding_forward(self, position_ids, **kw_args):
+        position = position_ids[..., : self.layout[1]]
+        position_plus = position_ids[..., self.layout[1] :] - self.original_sequence_length
+        position_embeddings = torch.cat(
+            (
+                self.transformer.position_embeddings(position),
+                self.get_mixin("extra_position_embedding").position_embeddings(position_plus),
+            ),
+            dim=-2,
+        )
+        return position_embeddings
+
+    def attention_forward(self, hidden_states, mask, layer_id=None, log_attention_weights=None, **kw_args):
+        attn_module = self.transformer.layers[layer_id].attention
+        # attention_plus on all layers
+        query_key_value_plus = self.get_mixin("attention_plus").query_key_value[layer_id]
+        dense_plus = self.get_mixin("attention_plus").dense[layer_id]
+        # split two parts
+        hidden_states_plus = hidden_states[:, self.layout[1] :]
+        hidden_states = hidden_states[:, : self.layout[1]]
+        # base model qkv
+        mixed_raw_layer = attn_module.query_key_value(hidden_states)
+        q0, k0, v0 = split_tensor_along_last_dim(mixed_raw_layer, 3)
+        # cuda2d model qkv
+        mixed_raw_layer = query_key_value_plus(hidden_states_plus)
+        q1, k1, v1 = split_tensor_along_last_dim(mixed_raw_layer, 3)
+
+        dropout_fn = attn_module.attention_dropout if self.training else None
+
+        # cuda2d attention
+        context_layer0, context_layer1 = sparse_attention_2d_light(
+            q0,
+            k0,
+            v0,
+            q1,
+            k1,
+            v1,
+            mask,
+            n_head=attn_module.num_attention_heads_per_partition,
+            text_len=self.layout[0],
+            kernel_size=self.kernel_size,
+            kernel_size2=self.kernel_size2,
+            attention_dropout=dropout_fn,
+            log_attention_weights=log_attention_weights,
+            add_scalar=(kw_args["add_scalar"] if "add_scalar" in kw_args else 0),
+        )
+
+        output_0 = attn_module.dense(context_layer0)
+        output_1 = dense_plus(context_layer1)
+        output = torch.cat((output_0, output_1), dim=1)
+
+        return output
+
+    def final_forward(self, logits, **kwargs):
+        logits_parallel = logits
+        logits_parallel = torch.nn.functional.linear(
+            logits_parallel.float(), self.transformer.word_embeddings.weight[:20000].float()
+        )
+        # logits_parallel = torch.nn.functional.linear(logits_parallel, self.transformer.word_embeddings.weight[:20000])
+        return logits_parallel
+
+    def disable_untrainable_params(self):
+        self.transformer.requires_grad_(False)
+
+    @classmethod
+    def add_model_specific_args(cls, parser):
+        group = parser.add_argument_group("Cuda2dModel", "cuda2d model configurations")
+        group.add_argument("--kernel-size", type=int, default=5)
+        group.add_argument("--kernel-size2", type=int, default=5)
+        group.add_argument("--layout", type=str, default="96,496,4096")
+        group.add_argument("--new-sequence-length", type=int, default=4096)
+        return parser
+
+
+def sparse_attention_2d_light(
+    q0,
+    k0,
+    v0,
+    q1,
+    k1,
+    v1,
+    attention_mask,
+    n_head,
+    text_len,
+    kernel_size=9,
+    kernel_size2=7,
+    attention_dropout=None,
+    log_attention_weights=None,
+    add_scalar=0,
+    **kwargs
+):
+    """
+    q0, k0, v0: [batch_size, 1088, hidden_size]
+    q1, k1, v1: [batch_size, 4096, h2]
+    n_head: int
+    attention_mask: [batch_size, 1088, 1088]
+    """
+    b, s0, h0 = q0.shape
+    b, s1, h1 = q1.shape
+    h, l0, l1 = h0 // n_head, sqrt(s0 - text_len), sqrt(s1)
+
+    q0 = q0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    v0 = v0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3)
+    k0T = k0.reshape(b, s0, n_head, h).permute(0, 2, 3, 1)
+
+    # standard attention for level 0
+    attention_scores = torch.matmul(q0 / math.sqrt(q0.shape[-1]), k0T)
+
+    if log_attention_weights is not None:
+        attention_scores += log_attention_weights
+    attention_scores = torch.mul(attention_scores, attention_mask) - 10000.0 * (1.0 - attention_mask)
+
+    attention_probs0 = F.softmax(attention_scores, dim=-1)
+
+    # local attention for level 1
+    q1 = (
+        (q1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1) / math.sqrt(h1 // n_head))
+        .contiguous()
+        .view(b * n_head, h1 // n_head, l1, l1)
+    )
+    k1 = k1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    v1 = v1.view(b, s1, n_head, h1 // n_head).permute(0, 2, 3, 1).contiguous().view(b * n_head, h1 // n_head, l1, l1)
+    # scores_1_to_1 = f_similar(q1, k1, kernel_size*2-1, kernel_size, True)
+    scores_1_to_1 = f_similar(q1, k1, kernel_size * 2 - 1, kernel_size, False)
+
+    # cross attention
+    k0T = k0T[..., -(l0**2) :].reshape(b * n_head, h, l0, l0).contiguous()
+    scores_1_to_0 = f_similar(q1, k0T, kernel_size2, kernel_size2, False)  # [b*n_head, l1, l1, field]
+    scores_1 = torch.cat(
+        (
+            scores_1_to_0.view(b * n_head, -1, scores_1_to_0.shape[3]) + add_scalar,
+            scores_1_to_1.view(b * n_head, -1, scores_1_to_1.shape[3]),
+        ),
+        dim=-1,
+    )
+    attention_probs1 = F.softmax(scores_1, dim=-1)
+
+    if attention_dropout is not None:
+        # with get_cuda_rng_tracker().fork():
+        attention_probs0 = attention_dropout(attention_probs0)
+        attention_probs1 = attention_dropout(attention_probs1)
+
+    # weighting for level 0
+    context0 = torch.matmul(attention_probs0, v0)  # [b, n_head, s0, h]
+    # weighting for level 1
+    probs_1_to_1 = attention_probs1[:, :, -scores_1_to_1.shape[3] :].view_as(scores_1_to_1)
+    # context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size*2-1, kernel_size, True)
+    context1_to_1 = f_weighting(v1, probs_1_to_1.contiguous(), kernel_size * 2 - 1, kernel_size, False)
+
+    context1 = context1_to_1.view(b, n_head * h, l1**2)
+    # weighting for cross attention
+    probs_1_to_0 = attention_probs1[:, :, : scores_1_to_0.shape[3]].view_as(scores_1_to_0)
+    v0_part = v0[:, :, -(l0**2) :].transpose(-1, -2).contiguous().view(b * n_head, h, l0, l0)
+    context1_to_0 = f_weighting(v0_part, probs_1_to_0.contiguous(), kernel_size2, kernel_size2, False)
+    context1_to_0 = context1_to_0.view(b, n_head * h, l1**2)
+    context1 = context1 + context1_to_0
+    return context0.transpose(1, 2).reshape(b, s0, h0), context1.transpose(-1, -2)