Follow along in this notebook.

Introduction

Fine-tuning an LLM involves taking a default model and tailoring it for a specific purpose or domain. This process enhances the model's performance and adaptability by training it on a smaller, more focused dataset relevant to the desired task. This additional training refines the model's capabilities and makes it more attuned to specific needs.

Fine-tuning is recommended when you want to teach the model a new task or leverage your company's unique knowledge base. Fine-tuning models is also helpful for generating a specific writing style or format.

For example, a global provider of video conferencing services needed to add call summarization and action item extraction features to its video conferencing offering in response to strong customer demand.

They built a solution by fine-tuning Cohere’s Command model on the nuances of transcript summarization. It accurately summarized live meetings in bullet points, including action items for call participants, and delivered it on the company’s proprietary conferencing platform.

Implementing an automated summarization system at a global provider of video conferencing services

Fine-tuning of Cohere models is available in both Amazon Bedrock and Amazon SageMaker. Sample notebooks on how to fine-tune a model on both platforms are available on the cohere-aws repository.

In this chapter, we’ll review one example using SageMaker to fine-tune the Command R model. At the time of writing, Command R fine-tuning is not yet available on Bedrock.

Code Walkthrough

We’ll go through the fine-tuning steps, followed by using a model once fine-tuning is complete.

Fine-Tuning

Let's install and import the necessary libraries and set up our Cohere client. Here we are using the cohere SDK together with AWS’s boto3 and sagemaker libraries.

# !pip install "cohere>=5.11.0"

import cohere
import boto3
import sagemaker as sage
from sagemaker.s3 import S3Uploader

The next step is to define the model package Amazon Resource Names (ARN) for the Command R model. The ARN is an identifying string for a SageMaker resource, and it varies between the regions where a resource is available.

Here, we define the Cohere package for the Command R model and map the model package against each region, which gives the complete ARN for each region.

region = boto3.Session().region_name

# Command R 082024 Finetuning
cohere_package = "cohere-command-r-v2-ft-cf30836984573101bba9b820364893bc"
# Legacy package for Command R Finetuning
# cohere_package = "cohere-command-r-ft-v-0-1-2-bae2282f0f4a30bca8bc6fea9efeb7ca"

# Mapping for algorithms
algorithm_map = {
    "us-east-1": f"arn:aws:sagemaker:us-east-1:865070037744:algorithm/{cohere_package}",
    "us-east-2": f"arn:aws:sagemaker:us-east-2:057799348421:algorithm/{cohere_package}",
    "us-west-2": f"arn:aws:sagemaker:us-west-2:594846645681:algorithm/{cohere_package}",
    "eu-central-1": f"arn:aws:sagemaker:eu-central-1:446921602837:algorithm/{cohere_package}",
    "ap-southeast-1": f"arn:aws:sagemaker:ap-southeast-1:192199979996:algorithm/{cohere_package}",
    "ap-southeast-2": f"arn:aws:sagemaker:ap-southeast-2:666831318237:algorithm/{cohere_package}",
    "ap-northeast-1": f"arn:aws:sagemaker:ap-northeast-1:977537786026:algorithm/{cohere_package}",
    "ap-south-1": f"arn:aws:sagemaker:ap-south-1:077584701553:algorithm/{cohere_package}",
}
if region not in algorithm_map.keys():
    raise Exception(f"Current boto3 session region {region} is not supported.")

arn = algorithm_map[region]

Next, we need to upload the dataset required for fine-tuning. The expected format for fine-tuning the Command R model via the Chat endpoint is a jsonl file, where each JSON object is a conversation containing a series of messages.

A System message in the beginning that acts as the preamble that guides the whole conversation
Multiple pairs of User and Chatbot messages representing the conversation that takes place between a human user and a chatbot

Here’s an example of what a data point looks like:

{'messages': 
 [{'role': 'System',
   'content': 'You are a chatbot trained to answer to my every question.'
  },
  {'role': 'User',
   'content': 'Hello'
  },
  {'role': 'Chatbot',
   'content': 'Greetings! How can I help you?'
  },
 {'role': 'User',
   'content': 'What makes a good running route?'
  },
  {'role': 'Chatbot',
   'content': 'A sidewalk-lined road is ideal so that you’re up and off the road away from vehicular traffic.'
  }
 ]
}

We’ll use the ScienceQA dataset for this demonstration, with the pre-formatted versions of the training and evaluation sets available. Here are two example data points from the dataset:

{"messages": [{"role": "User", "content": "Context: Fish live underwater. They have fins, not limbs. A tiger shark is an example of a fish.\nQuestion: Select the fish below.\nOptions:(A) black howler (B) cardinalfish\nAnswer:"}, {"role": "Chatbot", "content": "B"}]}

{"messages": [{"role": "User", "content": "Context: Select the best estimate.\nQuestion: How long is a guitar?\nOptions:(A) 1 centimeter (B) 1 meter (C) 1 kilometer (D) 1 millimeter\nAnswer:"}, {"role": "Chatbot", "content": "B"}]}

A few points to note:

The evaluation dataset is optional. If no evaluation dataset is provided, the training dataset will be auto-split into training and evaluation datasets at an 80:20 ratio.
The dataset must contain at least one example. If an evaluation dataset is absent, the training dataset must contain at least two examples.
However, we recommend using a dataset that contains at least 100 examples. Larger datasets are more likely to yield high-quality fine-tunes.

We can now upload the training and evaluation datasets to an S3 bucket using the S3Uploader.upload method.

s3_data_dir = "s3://..."  # Do not add a trailing slash otherwise the upload will not work

sess = sage.Session()

# TODO[Optional]: change it to your data
train_dataset = S3Uploader.upload("../data/scienceQA_train.jsonl", s3_data_dir, sagemaker_session=sess)

# optional eval dataset
eval_dataset = S3Uploader.upload("../data/scienceQA_eval.jsonl", s3_data_dir, sagemaker_session=sess)

print("traint_dataset", train_dataset)
print("eval_dataset", eval_dataset)

We also need to specify a directory on S3 where fine-tuned models should be stored.

s3_models_dir = f"s3://YOUR_S3_MODEL_PATH"  # DO NOT add a trailing slash at the end

Next, we create a Cohere client and define the set of hyperparameters to be used during fine-tuning. Hyperparameters are optional, but they provide more flexibility to developers who want to further optimize the fine-tuning process.

Developers can define the following hyperparameters:

train_epochs: The maximum number of training steps
learning_rate: The initial learning rate to be used during training
train_batch_size: The batch size during training
early_stopping_enabled: Enables early stopping
early_stopping_patience: Stops training if the loss metric does not improve beyond the early_stopping_threshold for this many rounds of evaluation
early_stopping_threshold: How much the loss must improve to prevent early stopping

# Example of how to pass hyperparameters to the fine-tuning job
train_parameters = {
    "train_epochs": 1,
    "early_stopping_patience": 2,
    "early_stopping_threshold": 0.001,
    "learning_rate": 0.01,
    "train_batch_size": 16,
}

We are now ready to start the fine-tuning process. To do this, we use the create_finetune method and specify the required arguments, as follows.

co = cohere.SagemakerClient(aws_region=region)

finetune_name = "test-finetune"

co.sagemaker_finetuning.create_finetune(arn=arn,
    name=finetune_name,
    train_data=train_dataset,
    eval_data=eval_dataset,
    s3_models_dir=s3_models_dir,
    instance_type="ml.p4de.24xlarge",
    training_parameters=train_parameters,
    role="ServiceRoleSagemaker",
)

The finetuned weights for the above will be store in a tar file {s3_models_dir}/test-finetune.tar.gz where the file name is the same as the name used during the creation of the finetune.

Inference

We can now use the fine-tuned model for inference. With SageMaker, we’ll need to create an endpoint via an AWS instance. The marketplace listing provides more details on the recommended instance type for a particular model, including pricing.

The steps are the same as the SageMaker example in an earlier chapter, with a few additional arguments. In particular, s3_models_dir defines the location where the fine-tuned model is stored.

endpoint_name="test-finetune"

co.sagemaker_finetuning.create_endpoint(arn=arn,
        endpoint_name=endpoint_name,
        s3_models_dir=s3_models_dir,
        recreate=True,
        instance_type="ml.p4de.24xlarge",
        role="ServiceRoleSagemaker",
)

# If the endpoint is already created, you just need to connect to it
co.sagemaker_finetuning.connect_to_endpoint(endpoint_name=endpoint_name)

We can then perform inference on the endpoint. For this, we call the Chat endpoint with the user message and get a response generated by the model.

message = "Select the best estimate.\nQuestion: How long is a guitar?\nOptions:(A) 32 feet (B) 32 yards (C) 32 inches (D) 32 miles\nAnswer:"

result = co.chat(message=message, return_prompt=True)
print(result)

cohere.Chat {
 response_id: ba6be7d5-0509-4038-adfb-7de695c264b0
 generation_id: 3b509b4f-b6b3-4eee-ae5d-ccb5e46f8822
 text: C
 chat_history: [{'role': 'USER', 'message': 'Select the best estimate.\nQuestion: How long is a guitar?\nOptions:(A) 32 feet (B) 32 yards (C) 32 inches (D) 32 miles\nAnswer:'}, {'role': 'CHATBOT', 'message': 'C'}]
 preamble: None
 finish_reason: COMPLETE
 token_count: None
 tool_calls: None
 citations: None
 documents: None
 search_results: None
 search_queries: None
 is_search_required: None
}

Optionally, we can test the fine-tuning performance on the evaluation dataset, as follows:

import json
from tqdm import tqdm
total = 0
correct = 0
for line in tqdm(open('../data/scienceQA_eval.jsonl').readlines()):
    total += 1
    question_answer_json = json.loads(line)
    question = question_answer_json["messages"][0]["content"]
    answer = question_answer_json["messages"][1]["content"]
    model_ans = co.sagemaker_finetuning.chat(message=question, temperature=0).text
    if model_ans == answer:
        correct +=1

print(f"Accuracy of finetuned model is %.3f" % (correct / total))

Important Note: You will continue to incur charges as long as an endpoint is running, so remember to delete it when your usage ends.

co.delete_endpoint()
co.close()

Conclusion

Fine-tuning an LLM is a powerful technique for customizing and enhancing its performance for specific tasks and domains. Businesses can use fine-tuning to improve accuracy and address unique scenarios specific to their domain. In this chapter, we explored how to fine-tune Cohere Command on Amazon SageMaker.

This chapter concludes our module, Cohere on AWS, and we hope you’ve found it useful and informative. We encourage you to continue your learning by exploring more modules at Cohere’s LLM University.

Contents

Introduction

Code Walkthrough

Fine-Tuning

Inference

Conclusion

Fine-Tuning Cohere Command R on Amazon SageMaker