code: def update_model_kwargs_fn(outputs: dict, **model_kwargs) -> dict:
      if "past_key_values" in outputs:
          model_kwargs["past"] = outputs["past_key_values"]
      else:
          model_kwargs["past"] = None
  
      # update token_type_ids with last value
      if "token_type_ids" in model_kwargs:
          token_type_ids = model_kwargs["token_type_ids"]
          model_kwargs["token_type_ids"] = torch.cat(
              [token_type_ids, token_type_ids[:, -1].unsqueeze(-1)],
              dim=-1,
          )
  
      # update attention mask
      if "attention_mask" in model_kwargs:
          attention_mask = model_kwargs["attention_mask"]
          model_kwargs["attention_mask"] = torch.cat(
              [
                  attention_mask,
                  attention_mask.new_ones((attention_mask.shape[0], 1)),
              ],
              dim=-1,
          )
  
      return model_kwargs
  
imports: import json
  import re
  from threading import Lock
  from typing import Any, Callable, Generator, List, Optional
  import jieba
  import torch
  import torch.distributed
  import torch.nn
  from pydantic import BaseModel, Field
label_desc: Use image node_id 4 for calling a global function with example usage: update_model_kwargs_fn(outputs) and returns: model_kwargs