Setting the file. One moment.
Nova Rlvr Reward Function Source Template · 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
Nova Rlvr Reward Function Source Template Back to the top of the page 10
Setup DevOps Agent
24.7
Code Output Guide · references
RDS Oracle
Reference
references/finetuning/templates/ nova_rlvr_reward_function_source_template.py
Python · 358 lines · 13 KB
# Add any other imports your use case requires
14
15 # ========================================================================================
16 # NOTE : INITIAL SUGGESTION ONLY - MUST BE CUSTOMIZED
17 #
18 # YOU MUST:
19 # 1. Review and update each section per YOUR use case
20 # 2. Customize the logic for YOUR SPECIFIC requirements
21 # 3. Replace example values (field names, thresholds, etc.) with your actual values
22 # 4. Test thoroughly before using
23 #
24 # DO NOT use this code as-is. It will not work until you uncomment and customize it.
25 # =========================================================================================
26
27
28 # =========================================================================================
29 # SECTION 1: Helper function — content normalization
30 # =========================================================================================
31 # Nova messages use content as a string, a list of {"type":"text","text":"..."} chunks,
32 # or a dict with a "text" key. This helper normalizes all forms to a plain string.
33 def content_to_text (content: Any) -> str :
34 """
35 Normalize Nova message content to a plain string.
36
37 Args:
38 content: String, list of text chunks, or dict with "text" key
39
40 Returns:
41 Plain text string
42 """
43 if content is None :
44 return ""
45 if isinstance (content, str ):
46 return content
47 if isinstance (content, list ):
48 parts: List[ str ] = []
49 for item in content:
50 if isinstance (item, str ):
51 parts.append(item)
52 elif isinstance (item, dict ) and "text" in item:
53 parts.append(item[ "text" ])
54 else :
55 parts.append( str (item))
56 return "" .join(parts)
57 if isinstance (content, dict ) and "text" in content:
58 return content[ "text" ]
59 return str (content)
60
61
62 # =========================================================================================
63 # SECTION 2: Helper function — ground truth extraction
64 # =========================================================================================
65 # Nova reference_answer can be a dict with flexible keys (answer, label, sentiment, etc.),
66 # a JSON string, or a plain string.
67 def coerce_ground_truth (ground_truth: Union[ str , Dict[ str , Any], Any]) -> Optional[ str ]:
68 """
69 Extract the ground-truth answer as a string from reference_answer.
70
71 Args:
72 ground_truth: Dict, JSON string, or plain string
73
74 Returns:
75 Ground truth string, or None if not found
76 """
77 if ground_truth is None :
78 return None
79
80 if isinstance (ground_truth, str ):
81 s = ground_truth.strip()
82 if not s:
83 return None
84 if s.startswith( "{" ) or s.startswith( "[" ):
85 try :
86 ground_truth = json.loads(s)
87 except Exception :
88 return s
89 else :
90 return s
91
92 if isinstance (ground_truth, dict ):
93 for key in ( "ground_truth" , "answer" , "label" , "sentiment" , "polarity" , "target" ):
94 if key in ground_truth and ground_truth[key] is not None :
95 return str (ground_truth[key])
96 if len (ground_truth) == 1 :
97 only_val = next ( iter (ground_truth.values()))
98 if only_val is not None :
99 return str (only_val)
100 return None
101
102 return str (ground_truth)
103
104
105 # =========================================================================================
106 # SECTION 3: Helper function — number extraction
107 # =========================================================================================
108 # TODO : UPDATE or REMOVE the helper function as per YOUR use case
109 # Note the below lines of code are examples and will not work for your use case
110 # You MUST update them to match YOUR use case
111 def extract_number (text: str ) -> Optional[ float ]:
112 """
113 Extract numerical answer from text.
114 Looks for numbers after answer keywords, or returns the last number found.
115
116 Args:
117 text: Text containing a numerical answer
118
119 Returns:
120 Extracted number as float, or None if no number found
121 """
122 if not text:
123 return None
124
125 # Try to find numbers after common answer keywords
126 answer_patterns = [
127 r " (?: equals | is | answer is | result is | = )\s * ( - ? \d + \. ? \d * ) " ,
128 r " (?: answer | result | solution ) : \s * ( - ? \d + \. ? \d * ) " ,
129 ]
130
131 for pattern in answer_patterns:
132 match = re.search(pattern, text, re. IGNORECASE )
133 if match:
134 try :
135 return float (match.group( 1 ))
136 except ValueError :
137 pass
138
139 # Fallback: find all numbers and return the last one (likely the answer)
140 pattern = r "- ? \d + \. ? \d * "
141 matches = re.findall(pattern, text)
142
143 if matches:
144 try :
145 return float (matches[ - 1 ])
146 except ValueError :
147 return None
148
149 return None
150
151
152 # =========================================================================================
153 # SECTION 4: Helper function — reasoning quality
154 # =========================================================================================
155 # TODO : UPDATE or REMOVE the helper function as per YOUR use case
156 # Note the below lines of code are examples and will not work for your use case
157 # You MUST update them to match YOUR use case
158 def compute_reasoning_quality (response: str ) -> float :
159 """
160 Compute reasoning quality score based on response characteristics.
161 This is a simple heuristic - customize based on your needs.
162
163 Args:
164 response: The model's response text
165
166 Returns:
167 Quality score between 0.0 and 1.0
168 """
169 if not response:
170 return 0.0
171
172 score = 0.0
173
174 # Check for reasoning indicators (customize these for your use case)
175 reasoning_indicators = [
176 "because" ,
177 "therefore" ,
178 "thus" ,
179 "since" ,
180 "so" ,
181 "first" ,
182 "second" ,
183 "then" ,
184 "finally" ,
185 "step" ,
186 "calculate" ,
187 "compute" ,
188 "equals" ,
189 ]
190
191 response_lower = response.lower()
192
193 # Award points for reasoning indicators (max 0.55)
194 indicator_count = sum ( 1 for indicator in reasoning_indicators if indicator in response_lower)
195 score += min (indicator_count * 0.11 , 0.55 )
196
197 # Award points for response length (indicates detailed reasoning, max 0.25)
198 if len (response) > 30 :
199 score += 0.05
200 if len (response) > 60 :
201 score += 0.1
202 if len (response) > 120 :
203 score += 0.1
204
205 # Award points for structured response (max 0.2)
206 if " \n " in response or "." in response:
207 score += 0.2
208
209 return min (score, 1.0 )
210
211
212 # =========================================================================================
213 # SECTION 5: Helper function — answer extraction
214 # =========================================================================================
215 # TODO : UPDATE or REMOVE the helper function as per YOUR use case
216 # Note the below lines of code are examples and will not work for your use case
217 # You MUST update them to match YOUR use case
218 def extract_answer (response: str ) -> Optional[ str ]:
219 """
220 Extract the answer from a Nova model response.
221 Looks for <|begin_of_solution|>...<|end_of_solution|> blocks and \\ boxed{} patterns.
222
223 Args:
224 response: The model's response text
225
226 Returns:
227 Extracted answer string, or None if not found
228 """
229 if not response:
230 return None
231
232 # Try solution block first
233 solution_match = re.search(
234 r "< \| begin_of_solution \| > (. *? ) < \| end_of_solution \| >" ,
235 response,
236 re. DOTALL ,
237 )
238 if solution_match:
239 boxed = re.findall( r " \\ boxed \{ ([ ^} ] + ) \} " , solution_match.group( 1 ))
240 if boxed:
241 return boxed[ - 1 ].strip()
242
243 # Fallback: boxed anywhere
244 boxed = re.findall( r " \\ boxed \{ ([ ^} ] + ) \} " , response)
245 if boxed:
246 return boxed[ - 1 ].strip()
247
248 return None
249
250
251 # =========================================================================================
252 # SECTION 6: Sample reward function
253 # =========================================================================================
254 # TODO : UPDATE or REMOVE the reward function as per YOUR use case
255 # Note the below lines of code are examples and will not work for your use case
256 # You MUST update them to match YOUR use case
257 def reward_function (sample: Dict[ str , Any], index: int ) -> Dict[ str , Any]:
258 """
259 Args:
260 sample: Dictionary containing messages and reference_answer
261 index: Sample index in batch
262
263 Returns:
264 Dictionary with reward scores and metrics
265 """
266 # ========================================================================
267 # SECTION 7: Parse input
268 # ========================================================================
269 # TODO : UPDATE logic to parse the input as per YOUR use case
270 # Note the below lines of code are examples and will not work for your use case
271 # You MUST update them to match YOUR use case
272 messages = sample.get( "messages" , [])
273 ground_truth = sample.get( "reference_answer" , {})
274
275 # Get the assistant's response (last message with role assistant or nova_assistant)
276 response = ""
277 for msg in messages:
278 role = msg.get( "role" , "" )
279 if role in ( "assistant" , "nova_assistant" ):
280 response = content_to_text(msg.get( "content" , "" ))
281
282 # Extract numerical answers
283 predicted = extract_number(response)
284 expected_str = coerce_ground_truth(ground_truth)
285 expected = extract_number(expected_str) if expected_str else None
286
287 # Compute metrics
288 exact_match = 0.0
289 answer_present = 0.0
290 reasoning_quality = compute_reasoning_quality(response)
291
292 if predicted is not None and expected is not None :
293 exact_match = 1.0 if abs (predicted - expected) < 1e-6 else 0.0
294 answer_present = 1.0
295
296 # ========================================================================
297 # SECTION 8: Compute reward scores
298 # ========================================================================
299 # TODO : UPDATE logic to compute aggregate score
300 # Note the below lines of code are examples and will not work for your use case
301 # You MUST update them to match YOUR use case
302 aggregate_reward = 0.7 * exact_match + 0.3 * reasoning_quality
303
304 # ========================================================================
305 # SECTION 9: Form the metrics list
306 # ========================================================================
307 # TODO : UPDATE logic to compute metrics list
308 # Note the below lines of code are examples and will not work for your use case
309 # You MUST update them to match YOUR use case
310 metrics = [
311 { "name" : "exact_match" , "value" : float (exact_match), "type" : "Reward" },
312 { "name" : "answer_present" , "value" : float (answer_present), "type" : "Metric" },
313 { "name" : "reasoning_quality" , "value" : float (reasoning_quality), "type" : "Metric" },
314 ]
315
316 # ========================================================================
317 # SECTION 10: Return output
318 # ========================================================================
319 # TODO : UPDATE the return statement to return YOUR output
320 # UPDATE the key before creating the evaluator
321 # Note the below lines of code are examples and will not work for your use case
322 # You MUST update them to match YOUR use case
323
324 return {
325 "id" : str (sample.get( "id" , f "sample- { index :03d} " )),
326 "aggregate_reward_score" : float (aggregate_reward),
327 "metrics_list" : metrics,
328 }
329
330
331 def lambda_handler (event: Dict[ str , Any], context: Any) -> Dict[ str , Any]:
332 """
333 AWS Lambda Handler for reward function.
334 SageMaker Nova RLVR invokes this with a bare list of samples and expects
335 a bare list of {id, aggregate_reward_score, ...} dicts in return.
336 """
337 # Event is a bare list of samples
338 batch = event if isinstance (event, list ) else []
339
340 results = []
341 for i, sample in enumerate (batch):
342 try :
343 result = reward_function(sample, i)
344 results.append(result)
345 except Exception as e:
346 results.append(
347 {
348 "id" : str (
349 sample.get( "id" , f "sample- { i :03d} " )
350 if isinstance (sample, dict )
351 else f "sample- { i :03d} "
352 ),
353 "aggregate_reward_score" : 0.0 ,
354 "metrics_list" : [],
355 }
356 )
357
358 return results