Deploy your finetuned model on AWS Marketplace

1
pip install "cohere>=5.11.0"

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 = (
6
    "cohere-command-r-v2-byoft-8370167e649c32a1a5f00267cd334c2c"
7
)
8
algorithm_map = {
9
    "us-east-1": f"arn:aws:sagemaker:us-east-1:865070037744:algorithm/{cohere_package}",
10
    "us-east-2": f"arn:aws:sagemaker:us-east-2:057799348421:algorithm/{cohere_package}",
11
    "us-west-2": f"arn:aws:sagemaker:us-west-2:594846645681:algorithm/{cohere_package}",
12
    "eu-central-1": f"arn:aws:sagemaker:eu-central-1:446921602837:algorithm/{cohere_package}",
13
    "ap-southeast-1": f"arn:aws:sagemaker:ap-southeast-1:192199979996:algorithm/{cohere_package}",
14
    "ap-southeast-2": f"arn:aws:sagemaker:ap-southeast-2:666831318237:algorithm/{cohere_package}",
15
    "ap-northeast-1": f"arn:aws:sagemaker:ap-northeast-1:977537786026:algorithm/{cohere_package}",
16
    "ap-south-1": f"arn:aws:sagemaker:ap-south-1:077584701553:algorithm/{cohere_package}",
17
}
18
if region not in algorithm_map:
19
    raise Exception(f"Current region {region} is not supported.")
20
arn = algorithm_map[region]
21
22
# The local directory of your adapter weights. No need to specify this, if you bring your own merged weights
23
adapter_weights_dir = "<adapter_weights_dir>"
24
25
# 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
26
merged_weights_dir = "<merged_weights_dir>"
27
28
# The S3 directory you want to save the merged weights
29
s3_checkpoint_dir = "<s3_checkpoint_dir>"
30
31
# 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
32
s3_output_dir = "<s3_output_dir>"
33
34
# The name of the export
35
export_name = "<export_name>"
36
37
# The name of the SageMaker endpoint
38
endpoint_name = "<endpoint_name>"
39
40
# The instance type for export and inference. Now "ml.p4de.24xlarge" and "ml.p5.48xlarge" are supported
41
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(
8
    base_model_name_or_path: str, adapter_weights_dir: str
9
):
10
    """
11
    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
12
    """
13
    base_model = CohereForCausalLM.from_pretrained(
14
        base_model_name_or_path
15
    )
16
    peft_model = PeftModel.from_pretrained(
17
        base_model, adapter_weights_dir
18
    )
19
    merged_model = peft_model.merge_and_unload()
20
    return merged_model
21
22
23
def save_hf_model(output_dir: str, model, tokenizer=None, args=None):
24
    """
25
    Save a HuggingFace model (and optionally tokenizer as well as additional args) to a local directory
26
    """
27
    os.makedirs(output_dir, exist_ok=True)
28
    model.save_pretrained(
29
        output_dir, state_dict=None, safe_serialization=True
30
    )
31
    if tokenizer is not None:
32
        tokenizer.save_pretrained(output_dir)
33
    if args is not None:
34
        torch.save(
35
            args, os.path.join(output_dir, "training_args.bin")
36
        )

1
# Get the merged model from adapter weights
2
merged_model = load_and_merge_model(
3
    "CohereForAI/c4ai-command-r-08-2024", adapter_weights_dir
4
)
5
6
# Save the merged weights to your local directory
7
save_hf_model(merged_weights_dir, merged_model)

1
sess = sage.Session()
2
merged_weights = S3Uploader.upload(
3
    merged_weights_dir, s3_checkpoint_dir, sagemaker_session=sess
4
)
5
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(
3
    endpoint_name=endpoint_name
4
)
5
6
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."
7
result = co.sagemaker_finetuning.chat(message=message)
8
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(
14
        message=question, temperature=0
15
    ).text
16
    if model_ans == answer:
17
        correct += 1
18
19
print(f"Accuracy of finetuned model is %.3f" % (correct / total))

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