mirror of https://github.com/hpcaitech/ColossalAI
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
77 lines
3.1 KiB
77 lines
3.1 KiB
1 year ago
|
import torch
|
||
|
|
import transformers
|
||
|
|
|
||
|
|
from ..registry import ModelAttribute, model_zoo
|
||
|
|
|
||
|
|
try:
|
||
|
|
from transformers import LlamaConfig, LlamaForCausalLM, LlamaForSequenceClassification, LlamaModel
|
||
|
|
HAS_LLAMA = True
|
||
|
|
except ImportError:
|
||
|
|
HAS_LLAMA = False
|
||
|
|
|
||
|
|
if HAS_LLAMA:
|
||
|
|
# ===============================
|
||
|
|
# Register LLaMA
|
||
|
|
# ===============================
|
||
|
|
|
||
|
|
def data_gen():
|
||
|
|
# the input ids are corresponding to the sentence
|
||
|
|
# 'Hello, my dog is cute'
|
||
|
|
#
|
||
|
|
# the code is give below:
|
||
|
|
# -----------------------------------
|
||
|
|
# from transformers import LlamaTokenizerFast
|
||
|
|
# tokenizer = LlamaTokenizerFast.from_pretrained("hf-internal-testing/llama-tokenizer")
|
||
|
|
# input = 'Hello, my dog is cute'
|
||
|
|
# tokenized_input = tokenizer(input, return_tensors='pt').to('cuda')
|
||
|
|
# -----------------------------------
|
||
|
|
|
||
|
|
input_ids = torch.Tensor([[1, 15043, 29892, 590, 11203, 338, 274, 1082]]).long()
|
||
|
|
attention_mask = torch.Tensor([[1, 1, 1, 1, 1, 1, 1, 1]]).long()
|
||
|
|
return dict(input_ids=input_ids, attention_mask=attention_mask)
|
||
|
|
|
||
|
|
# label is needed for casual lm
|
||
|
|
def data_gen_for_casual_lm():
|
||
|
|
data = data_gen()
|
||
|
|
labels = data['input_ids'].clone()
|
||
|
|
data['labels'] = labels
|
||
|
|
return data
|
||
|
|
|
||
|
|
# transform the output to a dict
|
||
|
|
output_transform_fn = lambda x: x
|
||
|
|
|
||
|
|
# function to get the loss
|
||
|
|
loss_fn = lambda output: output.last_hidden_state.mean()
|
||
|
|
loss_fn_for_casual_lm = lambda output: output.loss
|
||
|
|
loss_fn_for_seq_classification = lambda output: output.logits.mean()
|
||
|
|
|
||
|
|
config = LlamaConfig(num_hidden_layers=4,
|
||
|
|
hidden_size=128,
|
||
|
|
intermediate_size=256,
|
||
|
|
num_attention_heads=4,
|
||
|
|
max_position_embeddings=128,
|
||
|
|
num_labels=16)
|
||
|
|
|
||
|
|
# register the following models
|
||
|
|
# transformers.LlamaModel,
|
||
|
|
# transformers.LlamaForCausalLM,
|
||
|
|
# transformers.LlamaForSequenceClassification,
|
||
|
|
model_zoo.register(name='transformers_llama',
|
||
|
|
model_fn=lambda: transformers.LlamaModel(config),
|
||
|
|
data_gen_fn=data_gen,
|
||
|
|
output_transform_fn=output_transform_fn,
|
||
|
|
loss_fn=loss_fn,
|
||
|
|
model_attribute=ModelAttribute(has_control_flow=True))
|
||
|
|
model_zoo.register(name='transformers_llama_for_casual_lm',
|
||
|
|
model_fn=lambda: transformers.LlamaForCausalLM(config),
|
||
|
|
data_gen_fn=data_gen_for_casual_lm,
|
||
|
|
output_transform_fn=output_transform_fn,
|
||
|
|
loss_fn=loss_fn_for_casual_lm,
|
||
|
|
model_attribute=ModelAttribute(has_control_flow=True))
|
||
|
|
model_zoo.register(name='transformers_llama_for_sequence_classification',
|
||
|
|
model_fn=lambda: transformers.LlamaForSequenceClassification(config),
|
||
|
|
data_gen_fn=data_gen,
|
||
|
|
output_transform_fn=output_transform_fn,
|
||
|
|
loss_fn=loss_fn_for_seq_classification,
|
||
|
|
model_attribute=ModelAttribute(has_control_flow=True))
|