Fine-Tuning and Deploying Custom AI Models on Amazon Bedrock: A Practical Guide

In the rapidly evolving field of Generative AI, the ability to fine-tune and deploy custom models is a crucial skill that enables businesses to tailor solutions to their unique needs. which offers an excellent explanation.

In this article, I will guide you through the process of fine-tuning a language model using Amazon Bedrock. We'll focus on the most critical sections of the code, providing a clear understanding of the key components and steps involved in the fine-tuning process. The goal is to highlight the essential elements so you can grasp how the general workflow is implemented, without diving into every line of code.

For those who want to dive directly into the code or explore it further, the complete implementation is available in my GitHub repository.

Use Case: Summarizing Doctor-Patient Dialogues

For this example, we'll focus on a dataset containing doctor-patient dialogues sourced from the structure required for fine-tuning on Amazon Bedrock. Each line in the JSONL file must include a Prompt and a Completion field.

# Define output path for JSONL

output_file_name = 'clinical_notes_fine_tune.jsonl'

output_file_path = os.path.join('dataset', output_file_name)

output_dir = os.path.dirname(output_file_path)



# Prepare and save the dataset in the fine-tuning JSONL format

with open(output_file_path, 'w') as outfile:

    for _, row in train_dataset.iterrows():

        formatted_entry = {

            "completion": row['note'],  # Replace 'note' with the correct column name

            "prompt": f"Summarize the following conversation.\n\n{row['dialogue']}"  # Replace 'dialogue' as needed

        }

        json.dump(formatted_entry, outfile)

        outfile.write('\n')

    print(f"Dataset has been reformatted and saved to {output_file_path}.")

The following is one example of converted data into JSONL:

{

    "completion": "<Summarized clinical note>",

    "prompt": "Summarize the following conversation:\n\n<Doctor-patient dialogue>"

}

To make the dataset accessible for fine-tuning, it needs to be uploaded to an Amazon S3 bucket. The code ensures that the S3 bucket exists, creating it if necessary. Once the bucket is verified, the fine-tuning dataset, saved in JSON Lines format, is uploaded to the specified bucket. This step is essential, as Amazon Bedrock accesses the dataset from S3 during the fine-tuning process. The bucket name, region, and dataset file path are customizable, ensuring flexibility for various AWS configurations.

# Define the file path and S3 details

bucket_name = 'bedrock-finetuning-bucket25112024'

s3_key = abstracts_file



# Specify the region

region = 'us-east-1'  # Change this if needed



# Initialize S3 client with the specified region

s3_client = boto3.client('s3', region_name=region)



# Check if the bucket exists

try:

    existing_buckets = s3_client.list_buckets()

    bucket_exists = any(bucket['Name'] == bucket_name for bucket in existing_buckets['Buckets'])



    if not bucket_exists:

        # Create the bucket based on the region

        try:

            if bucket_region == 'us-east-1':

                # For us-east-1, do not specify LocationConstraint

                s3_client.create_bucket(Bucket=bucket_name)

                print(f"Bucket {bucket_name} created successfully in us-east-1.")

            else:

                # For other regions, specify the LocationConstraint

                s3_client.create_bucket(

                    Bucket=bucket_name,

                    CreateBucketConfiguration={'LocationConstraint': bucket_region}

                )

                print(f"Bucket {bucket_name} created successfully in {bucket_region}.")

        except Exception as e:

            print(f"Error creating bucket: {e}")

            raise e

    else:

        print(f"Bucket {bucket_name} already exists.")



    # Upload the file to S3

    try:

        s3_client.upload_file(output_file_path, bucket_name, s3_key)

        print(f"File uploaded to s3://{bucket_name}/{s3_key}")

    except Exception as e:

        print(f"Error uploading to S3: {e}")



except Exception as e:

    print(f"Error: {e}")

Step 3: Create and submit a fine-tuning job

With the dataset uploaded to Amazon S3 and the necessary resources in place, the next step is to create and submit the fine-tuning job. This involves specifying the pre-trained foundation model, the job details, and the fine-tuning parameters.

In this example, we fine-tune the Cohere command-light-text-v14 model to summarize medical conversations. Below is the configuration used to submit the job:

# Define the job parameters

base_model_id = "cohere.command-light-text-v14:7:4k"

job_name = "cohere-Summarizer-medical-finetuning-job-v1"

model_name = "cohere-Summarizer-medical-Tuned-v1"



# Submit the fine-tuning job

bedrock.create_model_customization_job(

    customizationType="FINE_TUNING",

    jobName=job_name,

    customModelName=model_name,

    roleArn=role_arn,

    baseModelIdentifier=base_model_id,

    hyperParameters={

        "epochCount": "1",  # Number of passes over the dataset

        "batchSize": "8",  # Number of samples per training step

        "learningRate": "0.00001",  # Learning rate for weight updates

    },

    trainingDataConfig={"s3Uri": f"s3://{bucket_name}/{s3_key}"},

    outputDataConfig={"s3Uri": f"s3://{bucket_name}/finetuned/"}

)

Key Parameters:

• 
  Base Model: The pre-trained model (cohere.command-light-text-v14) serves as the foundation for customization.


• 
  Job Name and Model Name: These identifiers help track the fine-tuning job and the resulting fine-tuned model for future deployments.

Hyperparameters:

• 
  epochCount: Specifies the number of training cycles. For demonstration, one epoch is used, but more epochs may yield better results for larger datasets.


• 
  batchSize: Determines how many samples are processed in each training step. A value of 8 balances memory usage and training efficiency.


• 
  learningRate: Sets the pace at which the model learns. Lower values ensure stable training but may require more time to converge.

Training and Output Configuration:The trainingDataConfig points to the S3 location of the dataset.The outputDataConfig specifies where the fine-tuned model will be stored.

Considerations:

The parameters, especially the hyperparameters, can be adjusted to optimize the fine-tuning process:

• 
  Smaller datasets may benefit from lower batchSize values.


• 
  Complex tasks may require more epochs to achieve convergence.


• 
  Learning rates should be fine-tuned to balance training stability and speed.

This step officially kicks off the fine-tuning process, allowing Amazon Bedrock to handle the heavy lifting of training your model with the provided data and configuration.

The status of the fine-tuning job can be also seen:

status = bedrock.get_model_customization_job(jobIdentifier="cohere-Summarizer-medical-finetuning-job-v1")["status"]

print(f"Job status: {status}")

The status of the fine-tuning job can be also seen in the Bedrock console:

Step 4: Test our fine-tuned model