claude-cookbooks/finetuning/finetuning_on_bedrock.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finetuning Claude 3 Haiku on Bedrock\n",
    "In this notebook, we'll walk you through the process of finetuning Claude 3 Haiku on Amazon Bedrock\n",
    "\n",
    "## What You'll Need\n",
    "- An AWS account with access to Bedrock\n",
    "- A dataset (or you can use the sample dataset provided here)\n",
    "- [A service role capable of accessing the s3 bucket where you save your training data](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-iam-role.html)\n",
    "\n",
    "## Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prep a Dataset\n",
    "Your dataset for bedrock finetuning needs to be a JSONL file (i.e. a file with a json object on each line).\n",
    "\n",
    "Each line in the JSONL file should be a JSON object with the following structure:\n",
    "\n",
    "```\n",
    "{\n",
    "  \"system\": \"<optional_system_message>\",\n",
    "  \"messages\": [\n",
    "    {\"role\": \"user\", \"content\": \"user message\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"assistant response\"},\n",
    "    ...\n",
    "  ]\n",
    "}\n",
    "```\n",
    "\n",
    "- The `system` field is optional.\n",
    "- There must be at least two messages.\n",
    "- The first message must be from the \"user\".\n",
    "- The last message must be from the \"assistant\".\n",
    "- User and assistant messages must alternate.\n",
    "- No extraneous keys are allowed.\n",
    "\n",
    "\n",
    "## Sample Dataset - JSON Mode\n",
    "We've included a sample dataset that teaches a model to respond to all questions with JSON. Here's what that dataset looks like:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "sample_dataset = []\n",
    "dataset_path = 'datasets/json_mode_dataset.jsonl'\n",
    "with open(dataset_path, 'r') as f:\n",
    "    for line in f:\n",
    "        sample_dataset.append(json.loads(line))\n",
    "\n",
    "print(json.dumps(sample_dataset[0], indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upload your dataset to S3\n",
    "Your dataset for finetuning should be available on s3; for this demo we'll write the sample dataset to an s3 bucket you control"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bucket_name = \"YOUR_BUCKET_NAME\"\n",
    "s3_path = \"json_mode_dataset.jsonl\"\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "s3.upload_file(dataset_path, bucket_name, s3_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch Bedrock Finetuning Job\n",
    "\n",
    "Now that you have your dataset ready, you can launch a finetuning job using `boto3`. First we'll configure a few parameters for the job:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration\n",
    "job_name = \"anthropic-finetuning-cookbook-training\"\n",
    "custom_model_name = \"anthropic_finetuning_cookbook\"\n",
    "role = \"YOUR_AWS_SERVICE_ROLE_ARN\"\n",
    "output_path = f\"s3://{bucket_name}/finetuning_example_results/\"\n",
    "base_model_id = \"arn:aws:bedrock:us-east-1::foundation-model/anthropic.claude-3-haiku-20240307-v1:0:200k\"\n",
    "\n",
    "# Hyperparameters\n",
    "epoch_count = 5\n",
    "batch_size = 4\n",
    "learning_rate_multiplier = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we can launch the job with `boto3`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock = boto3.client(service_name=\"bedrock\")\n",
    "bedrock_runtime = boto3.client(service_name=\"bedrock-runtime\")\n",
    "\n",
    "bedrock.create_model_customization_job(\n",
    "    customizationType=\"FINE_TUNING\",\n",
    "    jobName=job_name,\n",
    "    customModelName=custom_model_name,\n",
    "    roleArn=role,\n",
    "    baseModelIdentifier=base_model_id,\n",
    "    hyperParameters = {\n",
    "        \"epochCount\": f\"{epoch_count}\",\n",
    "        \"batchSize\": f\"{batch_size}\",\n",
    "        \"learningRateMultiplier\": f\"{learning_rate_multiplier}\",\n",
    "    },\n",
    "    trainingDataConfig={\"s3Uri\": f\"s3://{bucket_name}/{s3_path}\"},\n",
    "    outputDataConfig={\"s3Uri\": output_path},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use this to check the status of your job while its training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for the job status\n",
    "status = bedrock.get_model_customization_job(jobIdentifier=job_name)[\"status\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use your finetuned model!\n",
    "\n",
    "To use your finetuned model, [you'll need to host it using Provisioned Throughput in Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-use.html). Once your model is ready with Provisioned Throughput, you can invoked your model via the Bedrock API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "provisioned_throughput_arn = \"YOUR_PROVISIONED_THROUGHPUT_ARN\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock = boto3.client('bedrock-runtime', region_name = \"us-east-1\")\n",
    "body = json.dumps(\n",
    "    {\n",
    "        \"anthropic_version\": \"bedrock-2023-05-31\",\n",
    "        \"max_tokens\": 1000,\n",
    "        \"system\": \"JSON Mode: Enabled\",\n",
    "        \"messages\": [\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": [\n",
    "                    {\n",
    "                        \"type\": \"text\",\n",
    "                        \"text\":\"What is a large language model?\"\n",
    "                    }\n",
    "                ]\n",
    "            }\n",
    "        ]\n",
    "    }\n",
    ")\n",
    "response = bedrock_runtime.invoke_model(\n",
    "\tmodelId=provisioned_throughput_arn,\n",
    "    body=body\n",
    ")\n",
    "body = json.loads(response['body'].read().decode('utf-8'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(body['content'][0]['text'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}