code: def load_partial_heads(
      jiant_model,
      weights_dict,
      allow_missing_head_weights=False,
      allow_missing_head_model=False,
  ):
      """Loads model weights and returns lists of missing head weights or missing heads (if any).
  
      Args:
          jiant_model (JiantModel): jiant model (encoder and task models are core components).
          weights_dict (Dict): model weights.
          allow_missing_head_weights (bool): If False, throw exception if there are missing keys.
          allow_missing_head_model (bool): If False, throw exception if there are unexpected keys.
  
      Returns:
          Dict[str, List] containing lists of missing head weights or missing heads if any.
  
      """
      mismatch = jiant_model.load_state_dict(weights_dict, strict=False)
      result = {}
      if mismatch.missing_keys:
          assert allow_missing_head_weights
          missing_head_weights = set()
          for k in mismatch.missing_keys:
              missing_head_weights.add(k.split(".")[1])
          result["missing_head_weights"] = list(missing_head_weights)
      if mismatch.unexpected_keys:
          assert allow_missing_head_model
          missing_heads_model = set()
          for k in mismatch.unexpected_keys:
              missing_heads_model.add(k.split(".")[1])
          result["missing_heads_model"] = list(missing_heads_model)
      return result
  
imports: from dataclasses import dataclass
  from typing import Any
  from typing import Dict
  from typing import List
  import torch
  import torch.nn
  import transformers
  import warnings
  import jiant.proj.main.components.container_setup
  import jiant.proj.main.modeling.primary
  import jiant.utils.python.strings
  from jiant.proj.main.modeling.heads import JiantHeadFactory
  from jiant.proj.main.modeling.taskmodels import JiantTaskModelFactory, Taskmodel, MLMModel
  from jiant.shared.model_resolution import ModelArchitectures
  from jiant.tasks.core import Task
label_desc: Use image node_id 7 for calling a global function with example usage: load_partial_heads(jiant_model, weights_dict, allow_missing_head_weights, allow_missing_head_model) and returns: result