mirror of
https://github.com/wassname/Open-Assistant.git
synced 2026-06-30 16:40:05 +08:00
Sotirios Anagnostidis
32 lines
1.1 KiB
Python
32 lines
1.1 KiB
Python
from transformers import AutoModelForCausalLM
|
|
from .gptj import get_model as get_gptj_model
|
|
|
|
SUPPORTED_MODELS = ["galactica", "gpt-j"]
|
|
|
|
|
|
def freeze_top_n_layers(model, target_layers):
|
|
# its possible we can simply detect which module is a ModuleList
|
|
# and simply freeze the module without doing string parsing
|
|
for name, param in model.named_parameters():
|
|
if "embed" in name:
|
|
param.requires_grad = False
|
|
elif ".layer" in name or ".h." in name:
|
|
tokens = name.split(".")
|
|
layer_ = None
|
|
for token in tokens:
|
|
if token.isdigit():
|
|
layer_ = int(token)
|
|
break
|
|
|
|
if layer_ is not None and layer_ < target_layers:
|
|
# print('freeze ', layer_, name)
|
|
param.requires_grad = False
|
|
return model
|
|
|
|
|
|
def get_specific_model(model_name, cache_dir, quantization):
|
|
if "gpt-j" in model_name.lower():
|
|
return get_gptj_model(model_name, cache_dir, quantization)
|
|
else:
|
|
return AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
|