import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it does not exist
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None: # set to default bucket if a bucket name is not given sagemaker_session_bucket = sess.default_bucket() try: role = sagemaker.get_execution_role()
except ValueError: iam = boto3.client('iam') role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn'] sess = sagemaker.Session(default_bucket=sagemaker_session_bucket) print(f"sagemaker role arn: {role}")
print(f"sagemaker bucket: {sess.default_bucket()}")
print(f"sagemaker session region: {sess.boto_region_name}")

{ "id": "13818513", "summary": "Amanda baked cookies and will bring Jerry some tomorrow.", "dialogue": "Amanda: I baked cookies. Do you want some?rnJerry: Sure!rnAmanda: I'll bring you tomorrow :-)"
}

from transformers import AutoTokenizer model_id="bigscience/bloomz-7b1" # Load tokenizer of BLOOMZ
tokenized = AutoTokenizer.from_pretrained(model_id)
tokenizer.model_max_length = 2048 # overwrite wrong value

from random import randint
from itertools import chain
from functools import partial # custom instruct prompt start
prompt_template = f"Summarize the chat dialogue:n{{dialogue}}n---nSummary:n{{summary}}{{eos_token}}" # template dataset to add prompt to each sample
def template_dataset(sample): sample["text"] = prompt_template.format(dialogue=sample["dialogue"], summary=sample["summary"], eos_token=tokenizer.eos_token) return sample # apply prompt template per sample
dataset = dataset.map(template_dataset, remove_columns=list(dataset.features)) print(dataset[randint(0, len(dataset))]["text"]) # empty list to save remainder from batches to use in next batch
remainder = {"input_ids": [], "attention_mask": []} def chunk(sample, chunk_length=2048): # define global remainder variable to save remainder from batches to use in next batch global remainder # Concatenate all texts and add remainder from previous batch concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()} concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()} # get total number of tokens for batch batch_total_length = len(concatenated_examples[list(sample.keys())[0]]) # get max number of chunks for batch if batch_total_length >= chunk_length: batch_chunk_length = (batch_total_length // chunk_length) * chunk_length # Split by chunks of max_len. result = { k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)] for k, t in concatenated_examples.items() } # add remainder to global variable for next batch remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()} # prepare labels result["labels"] = result["input_ids"].copy() return result # tokenize and chunk dataset
lm_dataset = dataset.map( lambda sample: tokenizer(sample["text"]), batched=True, remove_columns=list(dataset.features)
).map( partial(chunk, chunk_length=2048), batched=True,
) # Print total number of samples
print(f"Total number of samples: {len(lm_dataset)}")

# save train_dataset to s3
training_input_path = f's3://{sess.default_bucket()}/processed/samsum-sagemaker/train'
lm_dataset.save_to_disk(training_input_path) print("uploaded data to:")
print(f"training dataset to: {training_input_path}") In [ ]:
training_input_path="s3://sagemaker-us-east-1-558105141721/processed/samsum-sagemaker/train"

def create_peft_config(model): from peft import ( get_peft_model, LoraConfig, TaskType, prepare_model_for_int8_training, ) peft_config = LoraConfig( task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, # Lora attention dimension. lora_alpha=32, # the alpha parameter for Lora scaling. lora_dropout=0.05, # the dropout probability for Lora layers. target_modules=["query_key_value"], ) # prepare int-8 model for training model = prepare_model_for_int8_training(model) model = get_peft_model(model, peft_config) model.print_trainable_parameters() return model

import time
# define Training Job Name job_name = f'huggingface-peft-{time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime())}' from sagemaker.huggingface import HuggingFace # hyperparameters, which are passed into the training job
hyperparameters ={ 'model_id': model_id, # pre-trained model 'dataset_path': '/opt/ml/input/data/training', # path where sagemaker will save training dataset 'epochs': 3, # number of training epochs 'per_device_train_batch_size': 1, # batch size for training 'lr': 2e-4, # learning rate used during training
} # create the Estimator
huggingface_estimator = HuggingFace( entry_point = 'run_clm.py', # train script source_dir = 'scripts', # directory which includes all the files needed for training instance_type = 'ml.g5.2xlarge', # instances type used for the training job instance_count = 1, # the number of instances used for training base_job_name = job_name, # the name of the training job role = role, # IAM role used in training job to access AWS resources, e.g. S3 volume_size = 300, # the size of the EBS volume in GB transformers_version = '4.26', # the transformers version used in the training job pytorch_version = '1.13', # the pytorch_version version used in the training job py_version = 'py39', # the python version used in the training job hyperparameters = hyperparameters
)

# define a data input dictionary with our uploaded s3 uris
data = {'training': training_input_path} # starting the train job with our uploaded datasets as inputs
huggingface_estimator.fit(data, wait=True)

from sagemaker.huggingface import HuggingFaceModel # 1. create Hugging Face Model Class
huggingface_model = HuggingFaceModel( model_data=huggingface_estimator.model_data, #model_data="s3://hf-sagemaker-inference/model.tar.gz", # Change to your model path role=role, transformers_version="4.26", pytorch_version="1.13", py_version="py39", model_server_workers=1
)

# 2. deploy model to SageMaker Inference
predictor = huggingface_model.deploy( initial_instance_count=1, instance_type= "ml.g5.4xlarge"
)

from random import randint
from datasets import load_dataset # 1. Load dataset from the hub
test_dataset = load_dataset("samsum", split="test") # 2. select a random test sample
sample = test_dataset[randint(0,len(test_dataset))] # 3. format the sample
prompt_template = f"Summarize the chat dialogue:n{{dialogue}}n---nSummary:n" fomatted_sample = { "inputs": prompt_template.format(dialogue=sample["dialogue"]), "parameters": { "do_sample": True, # sample output predicted probabilities "top_p": 0.9, # sampling technique Fan et. al (2018) "temperature": 0.1, # increasing the likelihood of high probability words and decreasing the likelihood of low probability words "max_new_tokens": 100, # }
} # 4. Invoke the SageMaker endpoint with the formatted sample
res = predictor.predict(fomatted_sample) # 5. Print the model output
print(res[0]["generated_text"].split("Summary:")[-1])
# Sample model output: Kirsten and Alex are going bowling this Friday at 7 pm. They will meet up and then go together.

print(sample["summary"])
# Sample model input: Kirsten reminds Alex that the youth group meets this Friday at 7 pm to go bowling.

predictor.delete_model()
predictor.delete_endpoint()