code: def load(model_name_or_path, heavy_hitter=False, args=None):
      print(f"Loading model from {model_name_or_path} ...")
      # however, tensor parallel for running falcon will occur bugs
      tokenizer = AutoTokenizer.from_pretrained(
          model_name_or_path,
          trust_remote_code=True,
      )
  
      if args is not None:
          config = AutoConfig.from_pretrained(model_name_or_path)
          config.hh_size = args.heavy_hitter_size
          config.recent_size = args.recent_size
  
      if heavy_hitter:
          model = H2OLlamaForCausalLM_streaming.from_pretrained(
              model_name_or_path,
              device_map="auto",
              config=config,
              torch_dtype=torch.float16,
              trust_remote_code=True,
          )
      else:
          model = AutoModelForCausalLM.from_pretrained(
              model_name_or_path,
              device_map="auto",
              torch_dtype=torch.float16,
              trust_remote_code=True,
          )
      if tokenizer.pad_token_id is None:
          if tokenizer.eos_token_id is not None:
              tokenizer.pad_token_id = tokenizer.eos_token_id
          else:
              tokenizer.pad_token_id = 0
  
      model.eval()
  
      return model, tokenizer
  
file_path: H2O/h2o_hf/utils_real_drop/stream.py
imports: import torch
  import argparse
  from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM
  import os.path
  import ssl
  import urllib.request
  import os
  import json
  from utils_real_drop.modify_llama import H2OLlamaForCausalLM_streaming
label_desc: Use image node_id 2 for calling a global function with example usage: load(model_name_or_path, heavy_hitter, args) and returns: model,  tokenizer