def get_adaptive_model(
lm_output_type,
prediction_heads,
layer_dims,
model,
device,
embeds_dropout_prob,
class_weights=None,
):
parsed_lm_output_types = lm_output_type.split(",")
language_model = LanguageModel.load(model)
initialized_heads = []
for head_name in prediction_heads.split(","):
initialized_heads.append(
PredictionHead.create(
prediction_head_name=head_name,
layer_dims=layer_dims,
class_weights=class_weights,
))
model = AdaptiveModel(
language_model=language_model,
prediction_heads=initialized_heads,
embeds_dropout_prob=embeds_dropout_prob,
lm_output_types=parsed_lm_output_types,
device=device,
)
return model
def get_adaptive_model(
lm_output_type,
prediction_heads,
layer_dims,
model,
device,
embeds_dropout_prob,
local_rank,
n_gpu,
fp16=False,
class_weights=None,
):
parsed_lm_output_types = lm_output_type.split(",")
initialized_heads = []
for head_name in prediction_heads.split(","):
initialized_heads.append(
PredictionHead.create(
prediction_head_name=head_name,
layer_dims=layer_dims,
class_weights=class_weights,
))
language_model = LanguageModel.load(model)
# TODO where are balance class weights?
model = AdaptiveModel(
language_model=language_model,
prediction_heads=initialized_heads,
embeds_dropout_prob=embeds_dropout_prob,
lm_output_types=parsed_lm_output_types,
device=device,
)
if fp16:
model.half()
if local_rank > -1:
model = WrappedDDP(model)
elif n_gpu > 1:
model = WrappedDataParallel(model)
return model