code: def load_model(model_name: str):
      print("Loading model...")
      tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
      model = transformers.AutoModelForCausalLM.from_pretrained(
          model_name, trust_remote_code=True, torch_dtype=torch.bfloat16
      )
  
      model = model.eval()
      model = model.to("cuda")
      model_wrapped = HF_LM(model, tokenizer, device="cuda")
      model_wrapped.o_proj_matrices = [
          torch.stack(p.weight.split(5120 // 40, dim=1)).detach()
          for nm, p in model_wrapped.model.named_modules()
          if "o_proj" in nm
      ]
      model_wrapped.linear_head = (
          model_wrapped.model.lm_head.weight.detach()
          .cpu()
          .float()
          .numpy()
      )
      return model_wrapped
  
imports: import streamlit
  import numpy
  import os
  import torch
  import pandas
  import json
  from tqdm import tqdm
  import transformers
  from model_lib.hf_tooling import HF_LM
  from easydict import EasyDict
  from scipy.special import softmax
  import seaborn
  import matplotlib.pyplot
  from hooks import *
  from model_lib.misc_utils import *
  from model_lib.streamlit_frontend import css
  from rare_knowledge.collection import *
label_desc: Use image node_id 6 for calling a global function with example usage: load_model(model_name) and returns: model_wrapped