Setting the file. One moment.
Rlaif Custom Prompt · AWS AI ML · aws/agent-toolkit-for-aws · Skills Docs
Issue No. 14 · AWS AI ML
↖ Back to the coverMessaging And Streaming Skills
Migration And Modernization Skills
Networking And Content Delivery Skills
Security And Identity Skills
Web And Mobile Development
120 chapters · 648 min
Back to the top of the page 10
Setup DevOps Agent
24.7
Code Output Guide · references
RDS Oracle
(opens in a new tab)
references/finetuning/code_templates/ rlaif_custom_prompt.py
Python · 182 lines · 6 KB
16 from sagemaker.ai_registry.air_constants import REWARD_PROMPT
17 from sagemaker.ai_registry.dataset import DataSet
18 from sagemaker.ai_registry.evaluator import Evaluator
19 from sagemaker.core import Attribution, set_attribution
20 from sagemaker.core.helper.session_helper import Session, get_execution_role
21 from sagemaker.core.resources import ModelPackageGroup
22 from sagemaker.train.common import TrainingType
23 from sagemaker.train.rlaif_trainer import RLAIFTrainer
24
25 set_attribution(Attribution. SAGEMAKER_AGENT_PLUGIN )
26
27 # Setup
28 sm_client = boto3.Session().client( "sagemaker" )
29 sagemaker_session = Session( sagemaker_client = sm_client)
30 bucket = sagemaker_session.default_bucket()
31
32 # Configuration - USER please fill in these fields with your information:
33
34 BASE_MODEL = "" # Sagemaker Hub model id
35 TRAINING_DATA_S3 = "" # S3 path
36 S3_OUTPUT_PATH = f "s3:// { bucket } /finetuning-output/"
37 ROLE_ARN = get_execution_role() # You can change this to a specific role
38 ACCEPT_EULA = (
39 False # Set to True to accept the base model's End-User License Agreement (OSS models only)
40 )
41 MODEL_PACKAGE_GROUP_NAME = "" # Auto-generated based on use case
42
43 # Reward model — the Bedrock LLM used as judge
44 # Available models and regions: see references/rlaif_guide.md
45 REWARD_MODEL_ID = ""
46
47 # Cell 3: Register Custom Reward Prompt
48
49 # Insert path to the custom Jinja prompt file (usually ../scripts/custom_reward_prompt.jinja)
50 CUSTOM_PROMPT_PATH = ""
51
52 reward_prompt_evaluator = Evaluator.create(
53 name = "[GENERATE A NAME FOR THE EVALUATOR HERE]" , # lowercase alphanumeric + hyphens, max 20 chars
54 type = REWARD_PROMPT ,
55 source = CUSTOM_PROMPT_PATH ,
56 sagemaker_session = sagemaker_session,
57 wait = True ,
58 )
59 REWARD_PROMPT_ARN = reward_prompt_evaluator.arn
60 print ( f "Reward Prompt Evaluator ARN: { REWARD_PROMPT_ARN } " )
61
62 # Cell 4: Create Dataset and Model Package Group
63
64 # Create Model Package Group
65 try :
66 model_package_group = ModelPackageGroup.create(
67 model_package_group_name = MODEL_PACKAGE_GROUP_NAME ,
68 model_package_group_description = "" ,
69 )
70 print ( f "Created new model package group named { MODEL_PACKAGE_GROUP_NAME } " )
71 except ClientError as e:
72 if e.response[ "Error" ][ "Code" ] in ( "ResourceInUse" , "ValidationException" ):
73 model_package_group = ModelPackageGroup.get(
74 model_package_group_name = MODEL_PACKAGE_GROUP_NAME
75 )
76 print (
77 f "There is already a model package group with the name { MODEL_PACKAGE_GROUP_NAME } . \n If you want to save your finetuned model under a different name, change the value of MODEL_PACKAGE_GROUP_NAME in the previous cell."
78 )
79 else :
80 raise
81
82 # Create Dataset
83 dataset = DataSet.create( name = MODEL_PACKAGE_GROUP_NAME , source = TRAINING_DATA_S3 , wait = True )
84 TRAINING_DATASET_ARN = dataset.arn
85
86 print ( f "Here is your model package group ARN: { model_package_group.model_package_group_arn }\n " )
87 print ( f "Here is your training dataset ARN: { dataset.arn } " )
88
89 # Cell 5: Configure Trainer
90
91 trainer = RLAIFTrainer(
92 model = BASE_MODEL ,
93 model_package_group = model_package_group,
94 reward_model_id = REWARD_MODEL_ID ,
95 reward_prompt = REWARD_PROMPT_ARN , # ARN of the registered custom prompt evaluator
96 training_dataset = TRAINING_DATASET_ARN ,
97 s3_output_path = S3_OUTPUT_PATH ,
98 sagemaker_session = sagemaker_session,
99 # accept_eula=ACCEPT_EULA, # Uncomment for Meta models
100 role = ROLE_ARN ,
101 )
102
103 print ( "Here are the recommended hyperparameters for the current training job:" )
104 print ( f "Batch size: { trainer.hyperparameters.global_batch_size } " )
105 print ( f "Learning rate: { trainer.hyperparameters.learning_rate } " )
106 print ( f "Epochs: { trainer.hyperparameters.max_epochs } " )
107
108 # Cell 6: Hyperparameter Overrides
109
110 # To change a hyperparameter, uncomment its corresponding line, and set the value you want.
111
112 # Note: If the value you choose is not supported for your model, you will get an error indicating the allowed range.
113
114 # Uncomment the following line to change the learning rate
115 # trainer.hyperparameters.learning_rate = 0.0002
116
117 # Uncomment the following line to change the batch size
118 # trainer.hyperparameters.global_batch_size = 16
119
120 # Uncomment the following line to change the number of epochs
121 # trainer.hyperparameters.max_epochs = 5
122
123 # Cell 7: Start Training
124
125 # Start training
126 training_job = trainer.train( wait = True )
127
128 print ( f "Training Job Name: { training_job.training_job_name } " )
129 print ( f "Training Status: { training_job.training_job_status } " )
130
131 # Save manifest
132 manifest_dir = Path( "[PROJECT_DIR]" ) / "manifests"
133 manifest_dir.mkdir( parents = True , exist_ok = True )
134 manifest_path = manifest_dir / f "training- { training_job.training_job_name } .json"
135 manifest_path.write_text(
136 json.dumps(
137 {
138 "training_job_name" : training_job.training_job_name,
139 "model_package_group_name" : MODEL_PACKAGE_GROUP_NAME ,
140 },
141 indent = 2 ,
142 )
143 )
144 print ( f "Manifest saved: { manifest_path } " )
145
146 # Cell 8: Plot and Display Metrics # NOTEBOOK_ONLY_SECTION
147
148 import matplotlib.pyplot as plt
149 import mlflow
150 from mlflow.tracking import MlflowClient
151
152 run_id = training_job.mlflow_details.mlflow_run_id
153 mlflow.set_tracking_uri(training_job.mlflow_config.mlflow_resource_arn)
154 client = MlflowClient()
155
156 metrics = [
157 "critic/rewards/mean" ,
158 "response_length/mean" ,
159 "actor/entropy_loss" ,
160 "actor/grad_norm" ,
161 "critic/advantages/mean" ,
162 ]
163
164 fig, axes = plt.subplots( 1 , len (metrics), figsize = ( 4 * len (metrics), 3 ))
165 for idx, metric in enumerate (metrics):
166 history = client.get_metric_history(run_id, metric)
167 if history:
168 axes[idx].plot(
169 [h.step for h in history],
170 [h.value for h in history],
171 linewidth = 2 ,
172 marker = "o" ,
173 markersize = 4 ,
174 )
175 axes[idx].set_xlabel( "Step" )
176 axes[idx].set_ylabel(metric.split( "/" )[ - 1 ])
177 axes[idx].set_title(metric, fontweight = "bold" )
178 axes[idx].grid( True , alpha = 0.3 )
179
180 plt.suptitle( f "Training Metrics: { training_job.training_job_name } " , fontweight = "bold" )
181 plt.tight_layout()
182 plt.show()