Amazon Bedrockで強化学習によるFine Tuning（RFT）

┌─────────────────────────────────────────────────────────────────┐
│  1. 報酬関数の実装 (Lambda)                                    　 │
│     モデル出力を評価し、スコア (0.0〜1.0) を返す                  　  │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│  2. トレーニングデータの準備 (JSONL)                                │
│     プロンプトのみ記載（正解は不要）                                  │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│  3. インフラのデプロイ (CDK)                                       │
│     Lambda, S3, IAM Role を作成                                  │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│  4. RFT ジョブの実行 (Bedrock API)                                │
│     Bedrock がモデル出力 → Lambda で評価 → 強化学習を繰り返す         │
└─────────────────────────────────────────────────────────────────┘
                              ↓
┌─────────────────────────────────────────────────────────────────┐
│  5. カスタムモデルのデプロイと評価                                   │
│     Deploy for on-demand → Playground or API でテスト            │
└─────────────────────────────────────────────────────────────────┘

{"messages": [{"role": "user", "content": "AIとは何か30文字以内で説明して"}]}
{"messages": [{"role": "user", "content": "クラウドコンピューティングを40文字以内で説明して"}]}
{"messages": [{"role": "user", "content": "機械学習とは何か25文字以内で答えて"}]}

[
  {
    "id": 0,
    "messages": [
      {"role": "user", "content": "AIとは何か30文字以内で説明して"},
      {"role": "assistant", "content": "人工知能の略で、機械が学習・推論する技術"}
    ]
  },
  {
    "id": 1,
    "messages": [
      {"role": "user", "content": "機械学習を25文字以内で説明して"},
      {"role": "assistant", "content": "データからパターンを学ぶ技術"}
    ]
  }
]

/// RFT から送られてくる入力
#[derive(Deserialize)]
struct RftSample {
    id: u64,                    // サンプルID（整数）
    messages: Vec<Message>,     // 会話履歴
}

/// Lambda が返す報酬スコア
#[derive(Serialize)]
struct RewardResult {
    id: u64,                        // 入力と同じID
    aggregate_reward_score: f64,    // 報酬スコア (0.0〜1.0)
}

fn extract_char_limit(user_message: &str) -> Option<usize> {
    let patterns = [
        r"(\d+)\s*文字以内",
        r"(\d+)\s*字以内",
    ];
    for pattern in patterns {
        if let Some(caps) = Regex::new(pattern).ok()?.captures(user_message) {
            return caps.get(1)?.as_str().parse().ok();
        }
    }
    None
}

fn calculate_reward(response: &str, limit: usize) -> f64 {
    let actual = response.trim().chars().count();  // Unicode文字数
    if actual <= limit { 1.0 } else { 0.0 }
}

/// ユーザーメッセージから文字数制限を抽出（例: "30文字以内で" → Some(30)）
fn extract_char_limit(user_message: &str) -> Option<usize> {
    let re = Regex::new(r"(\d+)\s*文字以内").ok()?;
    re.captures(user_message)?.get(1)?.as_str().parse().ok()
}

/// メッセージ配列からアシスタントの応答を取得
fn get_assistant_response(messages: &[Message]) -> Option<&str> {
    messages.iter()
        .filter(|m| m.role == "assistant")
        .last()
        .map(|m| m.content.as_str())
}

/// メッセージ配列からユーザーの最後のメッセージを取得
fn get_user_message(messages: &[Message]) -> Option<&str> {
    messages.iter()
        .filter(|m| m.role == "user")
        .last()
        .map(|m| m.content.as_str())
}

async fn handler(event: LambdaEvent<serde_json::Value>) -> Result<serde_json::Value, Error> {
    let samples: Vec<RftSample> = serde_json::from_value(event.payload)?;

    let results: Vec<RewardResult> = samples
        .into_iter()
        .map(|sample| {
            let user_msg = get_user_message(&sample.messages);
            let limit = user_msg.and_then(extract_char_limit);
            let response = get_assistant_response(&sample.messages);

            let score = match (response, limit) {
                (Some(res), Some(lim)) => calculate_reward(res, lim),
                _ => 0.0,  // 応答または制限がない場合は0点
            };

            RewardResult { id: sample.id, aggregate_reward_score: score }
        })
        .collect();

    Ok(serde_json::to_value(results)?)
}

$ cd lambda-reward
$ cargo lambda build --release --arm64
・・・

$ CDK_DEFAULT_ACCOUNT=xxxxx CDK_DEFAULT_REGION=us-east-1 \
% pnpm exec cdk deploy

# 出力:
# RftJsonRewardStack.RewardFunctionArn = arn:aws:lambda:us-east-1:xxxxx:function:rft-char-limit-reward
# RftJsonRewardStack.TrainingDataBucketName = xxxxxxx
# RftJsonRewardStack.BedrockRftRoleArn = arn:aws:iam::xxxxxx:role/BedrockRftJsonDemoRole

#サンプル例
{"messages": [{"role": "user", "content": "AIとは何か30文字以内で説明して"}]}
{"messages": [{"role": "user", "content": "クラウドコンピューティングを40文字以内で説明して"}]}
{"messages": [{"role": "user", "content": "機械学習とは何か25文字以内で答えて"}]}
{"messages": [{"role": "user", "content": "日本の首都を10文字以内で答えて"}]}

$ aws bedrock create-model-customization-job \
  --job-name "rft-char-limit" \
  --base-model-identifier "amazon.nova-2-lite-v1:0:256k" \
  --customization-type "REINFORCEMENT_FINE_TUNING" \
  --custom-model-name "nova-2-lite-char-limit-v1" \
  --role-arn "arn:aws:iam::xxxxxx:role/BedrockRftJsonDemoRole" \
  --training-data-config '{
    "s3Uri": "s3://<created bucket>/training-data/char-limit-training.jsonl"
  }' \
  --validation-data-config '{
    "validators": [{
      "s3Uri": "s3://<created bucket>/training-data/char-limit-validation.jsonl"
    }]
  }' \
  --output-data-config '{
    "s3Uri": "s3://<created bucket>/output/"
  }' \
  --customization-config '{
    "rftConfig": {
      "graderConfig": {
        "lambdaGrader": {
          "lambdaArn": "arn:aws:lambda:us-east-1:xxxxx:function:rft-char-limit-reward"
        }
      },
      "hyperParameters": {
        "epochCount": 10,
        "trainingSamplePerPrompt": 8
      }
    }
  }' \
  --region us-east-1

# カスタムモデルの ARN を取得
% aws bedrock list-custom-models --region us-east-1

# オンデマンド推論用にデプロイ
% aws bedrock create-custom-model-deployment \
    --model-deployment-name "rft-char-limit-deployment" \
    --model-id "<取得したカスタムモデルのARN>" \
    --region us-east-1

# ベースモデル
aws bedrock-runtime converse \
    --model-id "us.amazon.nova-2-lite-v1:0" \
    --messages '[{"role":"user","content":[{"text":"機械学習を20文字以内で説明して"}]}]' \
    --inference-config '{"temperature":0,"maxTokens":100}' \
    --region us-east-1

# カスタムモデル
aws bedrock-runtime converse \
    --model-id "arn:aws:bedrock:us-east-1:ACCOUNT_ID:custom-model-deployment/DEPLOYMENT_ID" \
    --messages '[{"role":"user","content":[{"text":"機械学習を20文字以内で説明して"}]}]' \
    --inference-config '{"temperature":0,"maxTokens":100}' \
    --region us-east-1

import * as cdk from "aws-cdk-lib";
import * as lambda from "aws-cdk-lib/aws-lambda";
import * as iam from "aws-cdk-lib/aws-iam";
import * as s3 from "aws-cdk-lib/aws-s3";
import * as s3deploy from "aws-cdk-lib/aws-s3-deployment";
import { Construct } from "constructs";
import * as path from "path";

export class RftRewardStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // S3 バケット（トレーニングデータ用）
    const bucket = new s3.Bucket(this, "TrainingDataBucket", {
      bucketName: `rft-demo-${this.account}-${this.region}`,
      removalPolicy: cdk.RemovalPolicy.DESTROY,
      autoDeleteObjects: true,
    });

    // トレーニングデータを S3 にデプロイ
    new s3deploy.BucketDeployment(this, "DeployTrainingData", {
      sources: [s3deploy.Source.asset(path.join(__dirname, "../../training-data"))],
      destinationBucket: bucket,
      destinationKeyPrefix: "training-data",
    });

    // Lambda 報酬関数（Rust）
    const rewardFunction = new lambda.Function(this, "RewardFunction", {
      functionName: "rft-char-limit-reward",
      runtime: lambda.Runtime.PROVIDED_AL2023,
      architecture: lambda.Architecture.ARM_64,
      handler: "bootstrap",
      code: lambda.Code.fromAsset(
        path.join(__dirname, "../../lambda-reward/target/lambda/rft-char-limit-reward")
      ),
      timeout: cdk.Duration.minutes(1),
      memorySize: 256,
    });

    // Bedrock からの Lambda 呼び出し許可
    rewardFunction.addPermission("BedrockInvoke", {
      principal: new iam.ServicePrincipal("bedrock.amazonaws.com"),
      action: "lambda:InvokeFunction",
    });

    // Bedrock RFT ジョブ用 IAM Role
    const bedrockRole = new iam.Role(this, "BedrockRftRole", {
      roleName: "BedrockRftRole",
      assumedBy: new iam.ServicePrincipal("bedrock.amazonaws.com"),
    });

    bucket.grantRead(bedrockRole);              // トレーニングデータ読み取り
    bucket.grantWrite(bedrockRole, "output/*"); // 出力書き込み
    rewardFunction.grantInvoke(bedrockRole);    // Lambda 呼び出し

    // 出力
    new cdk.CfnOutput(this, "LambdaArn", { value: rewardFunction.functionArn });
    new cdk.CfnOutput(this, "BucketName", { value: bucket.bucketName });
    new cdk.CfnOutput(this, "RoleArn", { value: bedrockRole.roleArn });
  }
}