Deploy Finetuned Command Models from AWS Marketplace

1
import cohere
2
import os
3
import sagemaker as sage
4
5
from sagemaker.s3 import S3Uploader

1
# Change "<aws_profile>" to your own AWS profile name
2
os.environ["AWS_PROFILE"] = "<aws_profile>"

1
# The AWS region
2
region = "<region>"
3
4
# Get the arn of the bring your own finetuning algorithm by region
5
cohere_package = "cohere-command-r-v2-byoft-8370167e649c32a1a5f00267cd334c2c"
6
algorithm_map = {
7
    "us-east-1": f"arn:aws:sagemaker:us-east-1:865070037744:algorithm/{cohere_package}",
8
    "us-east-2": f"arn:aws:sagemaker:us-east-2:057799348421:algorithm/{cohere_package}",
9
    "us-west-2": f"arn:aws:sagemaker:us-west-2:594846645681:algorithm/{cohere_package}",
10
    "eu-central-1": f"arn:aws:sagemaker:eu-central-1:446921602837:algorithm/{cohere_package}",
11
    "ap-southeast-1": f"arn:aws:sagemaker:ap-southeast-1:192199979996:algorithm/{cohere_package}",
12
    "ap-southeast-2": f"arn:aws:sagemaker:ap-southeast-2:666831318237:algorithm/{cohere_package}",
13
    "ap-northeast-1": f"arn:aws:sagemaker:ap-northeast-1:977537786026:algorithm/{cohere_package}",
14
    "ap-south-1": f"arn:aws:sagemaker:ap-south-1:077584701553:algorithm/{cohere_package}",
15
}
16
if region not in algorithm_map:
17
    raise Exception(f"Current region {region} is not supported.")
18
arn = algorithm_map[region]
19
20
# The local directory of your adapter weights. No need to specify this, if you bring your own merged weights
21
adapter_weights_dir = "<adapter_weights_dir>"
22
23
# The local directory you want to save the merged weights. Or the local directory of your own merged weights, if you bring your own merged weights
24
merged_weights_dir = "<merged_weights_dir>"
25
26
# The S3 directory you want to save the merged weights
27
s3_checkpoint_dir = "<s3_checkpoint_dir>"
28
29
# The S3 directory you want to save the exported TensorRT-LLM engine. Make sure you do not reuse the same S3 directory across multiple runs
30
s3_output_dir = "<s3_output_dir>"
31
32
# The name of the export
33
export_name = "<export_name>"
34
35
# The name of the SageMaker endpoint
36
endpoint_name = "<endpoint_name>"
37
38
# The instance type for export and inference. Now "ml.p4de.24xlarge" and "ml.p5.48xlarge" are supported
39
instance_type = "<instance_type>"

1
import torch
2
3
from peft import PeftModel
4
from transformers import CohereForCausalLM
5
6
7
def load_and_merge_model(base_model_name_or_path: str, adapter_weights_dir: str):
8
    """
9
    Load the base model and the model finetuned by PEFT, and merge the adapter weights to the base weights to get a model with merged weights
10
    """
11
    base_model = CohereForCausalLM.from_pretrained(base_model_name_or_path)
12
    peft_model = PeftModel.from_pretrained(base_model, adapter_weights_dir)
13
    merged_model = peft_model.merge_and_unload()
14
    return merged_model
15
16
17
def save_hf_model(output_dir: str, model, tokenizer=None, args=None):
18
    """
19
    Save a HuggingFace model (and optionally tokenizer as well as additional args) to a local directory
20
    """
21
    os.makedirs(output_dir, exist_ok=True)
22
    model.save_pretrained(output_dir, state_dict=None, safe_serialization=True)
23
    if tokenizer is not None:
24
        tokenizer.save_pretrained(output_dir)
25
    if args is not None:
26
        torch.save(args, os.path.join(output_dir, "training_args.bin"))
27
28
# Get the merged model from adapter weights
29
merged_model = load_and_merge_model("CohereForAI/c4ai-command-r-08-2024", adapter_weights_dir)
30
31
# Save the merged weights to your local directory
32
save_hf_model(merged_weights_dir, merged_model)

1
sess = sage.Session()
2
merged_weights = S3Uploader.upload(merged_weights_dir, s3_checkpoint_dir, sagemaker_session=sess)
3
print("merged_weights", merged_weights)

1
co = cohere.SagemakerClient(aws_region=region)
2
co.sagemaker_finetuning.export_finetune(
3
    arn=arn,
4
    name=export_name,
5
    s3_checkpoint_dir=s3_checkpoint_dir,
6
    s3_output_dir=s3_output_dir,
7
    instance_type=instance_type,
8
    role="ServiceRoleSagemaker",
9
)

1
co.sagemaker_finetuning.create_endpoint(
2
    arn=arn,
3
    endpoint_name=endpoint_name,
4
    s3_models_dir=s3_output_dir,
5
    recreate=True,
6
    instance_type=instance_type,
7
    role="ServiceRoleSagemaker",
8
)

1
# If the endpoint is already deployed, you can directly connect to it
2
co.sagemaker_finetuning.connect_to_endpoint(endpoint_name=endpoint_name)
3
4
message = "Classify the following text as either very negative, negative, neutral, positive or very positive: mr. deeds is , as comedy goes , very silly -- and in the best way."
5
result = co.sagemaker_finetuning.chat(message=message)
6
print(result)

1
import json
2
from tqdm import tqdm
3
4
eval_data_path = "<path_to_scienceQA_eval.jsonl>"
5
6
total = 0
7
correct = 0
8
for line in tqdm(open(eval_data_path).readlines()):
9
    total += 1
10
    question_answer_json = json.loads(line)
11
    question = question_answer_json["messages"][0]["content"]
12
    answer = question_answer_json["messages"][1]["content"]
13
    model_ans = co.sagemaker_finetuning.chat(message=question, temperature=0).text
14
    if model_ans == answer:
15
        correct += 1
16
17
print(f"Accuracy of finetuned model is %.3f" % (correct / total))

1
co.sagemaker_finetuning.delete_endpoint()
2
co.sagemaker_finetuning.close()