transformers_util module

unseal.transformers_util.get_num_layers(model: unseal.hooks.commons.HookedModel, layer_key_prefix: Optional[str] = None) → int

Get the number of layers in a model

Parameters

  • model (HookedModel) – The model to get the number of layers from

  • layer_key_prefix (Optional[str], optional) – The prefix to use for the layer keys, defaults to None

Returns

The number of layers in the model

Return type

int

unseal.transformers_util.load_from_pretrained(model_name: str, model_dir: Optional[str] = None, load_model: Optional[bool] = True, load_tokenizer: Optional[bool] = True, load_config: Optional[bool] = True, low_cpu_mem_usage: Optional[bool] = False) → Tuple[transformers.models.auto.modeling_auto.AutoModelForCausalLM, transformers.models.auto.tokenization_auto.AutoTokenizer, transformers.models.auto.configuration_auto.AutoConfig]

Load a pretrained model from huggingface’s transformer library

Parameters

  • model_name (str) – Name of the model, e.g. gpt2 or gpt-neo.

  • model_dir (Optional[str], optional) – Directory in which to look for the model, e.g. EleutherAI, defaults to None

  • load_model (Optional[bool], optional) – Whether to load the model itself, defaults to True

  • load_tokenizer (Optional[bool], optional) – Whether to load the tokenizer, defaults to True

  • load_config (Optional[bool], optional) – Whether to load the config file, defaults to True

  • low_cpu_mem_usage (bool, optional) – Whether to use low-memory mode, experimental feature of HF, defaults to False

Returns

model, tokenizer, config. Returns None values for those elements which were not loaded.

Return type

Tuple[AutoModelForCausalLM, AutoTokenizer, AutoConfig]