code: def read_in_k_v(state_dict, config):
      # for each of the encoder blocks:
      for i in range(config.num_encoder_blocks):
          for j in range(config.depths[i]):
              # read in weights + bias of keys and values (which is a single matrix in the original implementation)
              kv_weight = state_dict.pop(
                  f"segformer.encoder.block.{i}.{j}.attention.self.kv.weight"
              )
              kv_bias = state_dict.pop(
                  f"segformer.encoder.block.{i}.{j}.attention.self.kv.bias"
              )
              # next, add keys and values (in that order) to the state dict
              state_dict[
                  f"segformer.encoder.block.{i}.{j}.attention.self.key.weight"
              ] = kv_weight[: config.hidden_sizes[i], :]
              state_dict[
                  f"segformer.encoder.block.{i}.{j}.attention.self.key.bias"
              ] = kv_bias[: config.hidden_sizes[i]]
              state_dict[
                  f"segformer.encoder.block.{i}.{j}.attention.self.value.weight"
              ] = kv_weight[config.hidden_sizes[i] :, :]
              state_dict[
                  f"segformer.encoder.block.{i}.{j}.attention.self.value.bias"
              ] = kv_bias[config.hidden_sizes[i] :]
  
file_path: H2O/h2o_flexgen/benchmark/third_party/transformers/src/transformers/models/segformer/convert_segformer_original_to_pytorch.py
imports: import argparse
  import json
  from collections import OrderedDict
  from pathlib import Path
  import torch
  from PIL import Image
  import requests
  from huggingface_hub import hf_hub_download
  from transformers import SegformerConfig, SegformerFeatureExtractor, SegformerForImageClassification, SegformerForSemanticSegmentation
  from transformers.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: read_in_k_v(state_dict, config) without return types