1+ # Copyright 2024 Bytedance Ltd. and/or its affiliates
2+ #
3+ # Licensed under the Apache License, Version 2.0 (the "License");
4+ # you may not use this file except in compliance with the License.
5+ # You may obtain a copy of the License at
6+ #
7+ # http://www.apache.org/licenses/LICENSE-2.0
8+ #
9+ # Unless required by applicable law or agreed to in writing, software
10+ # distributed under the License is distributed on an "AS IS" BASIS,
11+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ # See the License for the specific language governing permissions and
13+ # limitations under the License.
14+
15+ import re
16+
17+
18+ def retrieve_format_reward (predict_str : str ) -> float :
19+ """
20+ Check if the prediction has the required format structure:
21+ <think>...</think> followed by <retrieval>...</retrieval> and <answer>...</answer>
22+
23+ Args:
24+ predict_str: The prediction string to evaluate
25+
26+ Returns:
27+ 1.0 if the format is correct, 0.0 otherwise
28+ """
29+ # Define the pattern for the required format
30+ pattern = re .compile (r"<think>.*</think>.*<retrieval>.*</retrieval>.*<answer>.*</answer>" , re .DOTALL )
31+ format_match = re .fullmatch (pattern , predict_str )
32+ return 1.0 if format_match else 0.0
33+
34+
35+ def extract_answer (predict_str : str ) -> str :
36+ """
37+ Extract the answer from the prediction string, looking for content within <answer> tags.
38+
39+ Args:
40+ predict_str: The prediction string to evaluate
41+
42+ Returns:
43+ The extracted answer or an empty string if not found
44+ """
45+ answer_match = re .search (r"<answer>(.*?)</answer>" , predict_str , re .DOTALL )
46+ if answer_match :
47+ return answer_match .group (1 ).strip ()
48+ return ""
49+
50+
51+ def retrieve_accuracy_reward (predict_str : str , ground_truth : str ) -> float :
52+ """
53+ Check if the answer (within <answer> tags) matches the ground truth.
54+
55+ Args:
56+ predict_str: The prediction string to evaluate
57+ ground_truth: The ground truth answer
58+
59+ Returns:
60+ 1.0 if the answer matches the ground truth, 0.0 otherwise
61+ """
62+ answer = extract_answer (predict_str )
63+
64+ # Normalize both answers for comparison (lowercase, strip spaces)
65+ answer_norm = answer .lower ().strip ()
66+ ground_truth_norm = ground_truth .lower ().strip ()
67+
68+ # Check if the normalized answer matches the ground truth
69+ return 1.0 if answer_norm == ground_truth_norm else 0.0
70+
71+
72+ def retrieval_spans_in_context (predict_str : str , context : str ) -> float :
73+ """
74+ Check if all retrieval spans in the prediction are found in the context.
75+
76+ Args:
77+ predict_str: The prediction string to evaluate
78+ context: The context string to search in
79+
80+ Returns:
81+ 1.0 if all retrieval spans are in the context, 0.0 otherwise
82+ """
83+ # Extract all retrieval spans
84+ spans = re .findall (r"<retrieval>(.*?)</retrieval>" , predict_str , re .DOTALL )
85+
86+ # If no retrieval spans were found, return 0.0
87+ if not spans :
88+ return 0.0
89+
90+ # Check if all spans are in the context
91+ spans_found = 0
92+ for span in spans :
93+ # Clean up the span by removing extra whitespace
94+ cleaned_span = re .sub (r'\s+' , ' ' , span ).strip ()
95+ if not cleaned_span :
96+ continue
97+ if cleaned_span in context :
98+ spans_found += 1
99+
100+ # Return a score based on the proportion of spans found
101+ if not spans :
102+ return 0.0
103+ return min (1.0 , spans_found / len ([s for s in spans if s .strip ()]))
104+
105+
106+ def retrieve_compute_score (predict_str : str , ground_truth : str , context : str ) -> float :
107+ """
108+ Compute the combined score for retrieval-based QA evaluation.
109+
110+ Args:
111+ predict_str: The prediction string to evaluate
112+ ground_truth: The ground truth answer
113+ context: The context from which retrieval should happen
114+
115+ Returns:
116+ The combined reward score between 0.0 and 1.0
117+ """
118+ # Calculate individual reward components
119+ format_score = retrieve_format_reward (predict_str )
120+ accuracy_score = retrieve_accuracy_reward (predict_str , ground_truth )
121+ retrieval_score = retrieval_spans_in_context (predict_str , context )
122+
123+ # Combine scores with weights (similar to math.py's weighting)
124+ return 0.7 * accuracy_score + 0.1 * format_score + 0.2 * retrieval_score
0 commit comments