feat: refactor whole code base, add logic for training R1 distil base models, change some template and reward logics
- Break down rl_helpers into smaller modules - Removed deprecated rl_helpers module to streamline the codebase. - Enhance initial user prompt template inspired by Search-R1main
thinhlpg
1 month ago
parent
c90c03267e
commit
31dcbf5d8a
@ -0,0 +1,43 @@
|
|||||||
|
"""
|
||||||
|
Main package exports for RL helpers.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from trl.trainer.grpo_trainer import apply_chat_template
|
||||||
|
|
||||||
|
from src.agent import Agent
|
||||||
|
from src.config import logger
|
||||||
|
from src.evaluation import check_student_answers, run_eval, verify
|
||||||
|
from src.prompts import build_user_prompt, format_search_results, get_system_prompt
|
||||||
|
from src.rewards import (
|
||||||
|
build_reward_correctness_fn,
|
||||||
|
reward_em_chunk,
|
||||||
|
reward_format,
|
||||||
|
reward_retry,
|
||||||
|
)
|
||||||
|
from src.search_module import get_qa_dataset, search
|
||||||
|
from src.tokenizer_adapter import LlamaTokenizerAdapter, R1DistilTokenizerAdapter
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
# Prompts
|
||||||
|
"get_system_prompt",
|
||||||
|
"build_user_prompt",
|
||||||
|
"format_search_results",
|
||||||
|
"apply_chat_template",
|
||||||
|
# Agent
|
||||||
|
"Agent",
|
||||||
|
"LlamaTokenizerAdapter",
|
||||||
|
"R1DistilTokenizerAdapter",
|
||||||
|
# Rewards
|
||||||
|
"build_reward_correctness_fn",
|
||||||
|
"reward_format",
|
||||||
|
"reward_retry",
|
||||||
|
"reward_em_chunk",
|
||||||
|
# Evaluation
|
||||||
|
"run_eval",
|
||||||
|
"check_student_answers",
|
||||||
|
"verify",
|
||||||
|
# Search
|
||||||
|
"get_qa_dataset",
|
||||||
|
"search",
|
||||||
|
"logger",
|
||||||
|
]
|
||||||
@ -0,0 +1,238 @@
|
|||||||
|
"""
|
||||||
|
Evaluation utilities for RL training.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import inspect
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
from src.agent import Agent
|
||||||
|
from src.config import logger
|
||||||
|
from src.search_module import get_qa_dataset
|
||||||
|
from src.tokenizer_adapter import LlamaTokenizerAdapter, R1DistilTokenizerAdapter
|
||||||
|
|
||||||
|
|
||||||
|
async def verify(student_answer: str, question: str, answer: str) -> bool:
|
||||||
|
"""
|
||||||
|
Verify if student's answer matches the correct answer.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
student_answer: The model's answer
|
||||||
|
question: The original question
|
||||||
|
answer: The ground truth answer
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
bool: True if answer is correct, False otherwise
|
||||||
|
"""
|
||||||
|
logger.debug(f"Verifying answer for question: {question}")
|
||||||
|
logger.debug(f"Student answer: {student_answer}")
|
||||||
|
logger.debug(f"Correct answer: {answer}")
|
||||||
|
|
||||||
|
# Simple string matching for now
|
||||||
|
# TODO: Implement more sophisticated matching
|
||||||
|
return student_answer.strip().lower() == answer.strip().lower()
|
||||||
|
|
||||||
|
|
||||||
|
def check_student_answers(
|
||||||
|
questions: list[str],
|
||||||
|
answers: list[str],
|
||||||
|
student_answers: list, # Can be strings or dicts
|
||||||
|
vllm_generate_func,
|
||||||
|
tokenizer,
|
||||||
|
log_file=None,
|
||||||
|
) -> list[bool]:
|
||||||
|
"""
|
||||||
|
Evaluates a list of student answers against the true answers using a vLLM generate function.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
questions: List of questions
|
||||||
|
answers: List of correct answers
|
||||||
|
student_answers: List of student answers to evaluate
|
||||||
|
vllm_generate_func: Function to generate verification responses
|
||||||
|
tokenizer: Tokenizer for formatting prompts
|
||||||
|
log_file: Optional path to write detailed results
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of boolean results (True for correct answers)
|
||||||
|
"""
|
||||||
|
logger.info(f"Checking {len(questions)} student answers")
|
||||||
|
|
||||||
|
if not (len(questions) == len(answers) == len(student_answers)):
|
||||||
|
logger.error("Mismatched lengths between questions, answers, and student answers")
|
||||||
|
raise ValueError("The number of questions, answers, and student answers must be equal.")
|
||||||
|
|
||||||
|
prompts = []
|
||||||
|
for question, answer, student_ans in zip(questions, answers, student_answers):
|
||||||
|
prompt_text = (
|
||||||
|
"You are grading a student's answer to a question. For the following question, "
|
||||||
|
"compare the student's answer to the correct answer. Reply with 'Yes' if the student's answer contains the correct information, "
|
||||||
|
"even if it's not an exact match. If the student's answer doesn't contain the right information or is completely incorrect, reply with 'No'.\n\n"
|
||||||
|
f"Question: {question}\n"
|
||||||
|
f"Correct Answer: {answer}\n"
|
||||||
|
f"Student Answer: {student_ans}\n\n"
|
||||||
|
"Your response should be just 'Yes' or 'No'."
|
||||||
|
)
|
||||||
|
|
||||||
|
formatted_prompt = tokenizer.apply_chat_template(
|
||||||
|
[{"role": "user", "content": prompt_text}],
|
||||||
|
tokenize=False,
|
||||||
|
add_generation_prompt=True,
|
||||||
|
)
|
||||||
|
prompts.append(formatted_prompt)
|
||||||
|
logger.debug(f"Created verification prompt for question: {question[:50]}...")
|
||||||
|
|
||||||
|
logger.info("Generating verification responses")
|
||||||
|
responses = vllm_generate_func(prompts)
|
||||||
|
responses_text = []
|
||||||
|
for response in responses:
|
||||||
|
# Handle different response formats
|
||||||
|
if hasattr(response, "outputs"):
|
||||||
|
try:
|
||||||
|
responses_text.append(response.outputs[0].text)
|
||||||
|
except (AttributeError, IndexError):
|
||||||
|
# Fallback for simple string responses
|
||||||
|
responses_text.append(str(response))
|
||||||
|
else:
|
||||||
|
responses_text.append(str(response))
|
||||||
|
logger.debug(f"Got {len(responses_text)} verification responses")
|
||||||
|
|
||||||
|
results = []
|
||||||
|
for response in responses_text:
|
||||||
|
results.append("yes" in response.lower())
|
||||||
|
logger.debug(f"Verification result: {'yes' in response.lower()}")
|
||||||
|
|
||||||
|
logger.info(f"Verification complete. {sum(results)}/{len(results)} answers correct")
|
||||||
|
|
||||||
|
# Append the QA details and verifier's response to the specified log file
|
||||||
|
if log_file:
|
||||||
|
with open(log_file, "a") as file:
|
||||||
|
timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||||
|
file.write(f"\n📝 === QA Evaluation at {timestamp} ===\n")
|
||||||
|
file.write(f"📂 File: {__file__}\n")
|
||||||
|
|
||||||
|
# Get current frame info safely
|
||||||
|
frame = inspect.currentframe()
|
||||||
|
if frame:
|
||||||
|
file.write(f"📍 Line: {frame.f_lineno}\n")
|
||||||
|
# Don't forget to delete the frame to avoid reference cycles
|
||||||
|
del frame
|
||||||
|
|
||||||
|
file.write("=" * 80 + "\n")
|
||||||
|
|
||||||
|
for i, (question, answer, student_ans, verifier_response) in enumerate(
|
||||||
|
zip(questions, answers, student_answers, responses_text)
|
||||||
|
):
|
||||||
|
file.write(f"\n❓ Question {i + 1}:\n")
|
||||||
|
file.write("-" * 40 + "\n")
|
||||||
|
file.write(f"📋 Question: {question}\n")
|
||||||
|
file.write(f"✅ Correct Answer: {answer}\n")
|
||||||
|
file.write(f"👨🎓 Student Answer: {student_ans}\n")
|
||||||
|
file.write(f"🔍 Verifier said: {verifier_response}\n")
|
||||||
|
|
||||||
|
# Add search results if available in the chat state
|
||||||
|
if isinstance(student_ans, dict) and "messages" in student_ans:
|
||||||
|
# Get messages from dict
|
||||||
|
messages = student_ans.get("messages", [])
|
||||||
|
search_results = [msg.get("content", "") for msg in messages if msg.get("role") == "ipython"]
|
||||||
|
if search_results:
|
||||||
|
file.write("\n🔎 Search Results:\n")
|
||||||
|
for j, result in enumerate(search_results, 1):
|
||||||
|
file.write(f"\nSearch {j}:\n{result}\n")
|
||||||
|
|
||||||
|
file.write("-" * 40 + "\n")
|
||||||
|
|
||||||
|
file.write(
|
||||||
|
f"\n📊 Summary: {sum(results)}/{len(results)} answers correct ({sum(results) / len(results) * 100:.2f}%)\n"
|
||||||
|
)
|
||||||
|
file.write("=" * 80 + "\n\n")
|
||||||
|
|
||||||
|
return results
|
||||||
|
|
||||||
|
|
||||||
|
def run_eval(generate_fn, verify_fn, tokenizer, output_file=None, debug_file=None):
|
||||||
|
"""
|
||||||
|
Run evaluation on the test dataset and return results.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
generate_fn: Function to generate completions
|
||||||
|
verify_fn: Function to verify results
|
||||||
|
tokenizer: Tokenizer for processing text
|
||||||
|
output_file: Path to save evaluation results summary
|
||||||
|
debug_file: Path to save detailed debug information
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
full_chat_states: The chat states from evaluation
|
||||||
|
"""
|
||||||
|
train_dataset, test_dataset = get_qa_dataset()
|
||||||
|
questions = test_dataset["prompt"]
|
||||||
|
|
||||||
|
# Create agent with appropriate adapter based on tokenizer
|
||||||
|
tokenizer_name = tokenizer.name_or_path.lower()
|
||||||
|
if "deepseek-r1-distill" in tokenizer_name:
|
||||||
|
adapter = R1DistilTokenizerAdapter()
|
||||||
|
elif "llama" in tokenizer_name:
|
||||||
|
adapter = LlamaTokenizerAdapter()
|
||||||
|
else:
|
||||||
|
adapter = R1DistilTokenizerAdapter()
|
||||||
|
|
||||||
|
agent = Agent(adapter)
|
||||||
|
agentic_outputs = agent.run_agent(generate_fn, tokenizer, questions)
|
||||||
|
full_chat_states = agentic_outputs.full_chat_states
|
||||||
|
final_responses = agentic_outputs.final_response_str
|
||||||
|
rewards = verify_fn(questions, full_chat_states, answer=test_dataset["answer"])
|
||||||
|
|
||||||
|
# Calculate results
|
||||||
|
percent_correct = sum(rewards) / len(rewards) * 100
|
||||||
|
|
||||||
|
# Log results to console
|
||||||
|
logger.info("RESULTS:")
|
||||||
|
logger.info(f"percentage of correct answers: {percent_correct:.2f}%")
|
||||||
|
logger.info("=" * 30)
|
||||||
|
|
||||||
|
# Save results to file if specified
|
||||||
|
if output_file:
|
||||||
|
try:
|
||||||
|
with open(output_file, "w") as f:
|
||||||
|
f.write("EVALUATION RESULTS\n")
|
||||||
|
f.write("=================\n\n")
|
||||||
|
f.write(f"Total questions: {len(questions)}\n")
|
||||||
|
f.write(f"Correct answers: {sum(rewards)}\n")
|
||||||
|
f.write(f"Percentage correct: {percent_correct:.2f}%\n\n")
|
||||||
|
|
||||||
|
f.write("Individual results:\n")
|
||||||
|
for i, (q, r, resp) in enumerate(zip(questions, rewards, final_responses)):
|
||||||
|
f.write(f"\nQ{i + 1}: {q[:100]}...\n")
|
||||||
|
f.write(f"Correct: {'✓' if r else '✗'}\n")
|
||||||
|
f.write(f"Response: {resp[:150]}...\n")
|
||||||
|
f.write("-" * 40 + "\n")
|
||||||
|
logger.info(f"Saved evaluation results to {output_file}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error saving results file: {e}")
|
||||||
|
|
||||||
|
# Save debug information if specified
|
||||||
|
if debug_file:
|
||||||
|
try:
|
||||||
|
import json
|
||||||
|
|
||||||
|
debug_data = []
|
||||||
|
for i, (q, r, resp, chat) in enumerate(zip(questions, rewards, final_responses, full_chat_states)):
|
||||||
|
debug_data.append(
|
||||||
|
{
|
||||||
|
"question_id": i,
|
||||||
|
"question": q,
|
||||||
|
"is_correct": bool(r),
|
||||||
|
"final_response": resp,
|
||||||
|
"chat_state": {
|
||||||
|
k: str(v) if isinstance(v, (list, dict)) else v
|
||||||
|
for k, v in chat.items()
|
||||||
|
if k != "tokenizer"
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
with open(debug_file, "w") as f:
|
||||||
|
json.dump(debug_data, f, indent=2)
|
||||||
|
logger.info(f"Saved debug information to {debug_file}")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error saving debug file: {e}")
|
||||||
|
|
||||||
|
return full_chat_states
|
||||||
@ -0,0 +1,67 @@
|
|||||||
|
"""
|
||||||
|
Prompt-related functions for handling system and user prompts.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
|
|
||||||
|
|
||||||
|
def get_system_prompt():
|
||||||
|
"""Get the system prompt with current date."""
|
||||||
|
current_date = datetime.now().strftime("%d %b %Y")
|
||||||
|
return f"""Cutting Knowledge Date: December 2023
|
||||||
|
Today Date: {current_date}
|
||||||
|
|
||||||
|
You are a helpful assistant with search capabilities.
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def build_user_prompt(q):
|
||||||
|
"""
|
||||||
|
Build a user prompt with the question using the new template format.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
q (str): The question to ask
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
str: Formatted user prompt
|
||||||
|
"""
|
||||||
|
user_prompt = f"""Answer the given question. \
|
||||||
|
You must conduct reasoning inside <think> and </think> first every time you get new information. \
|
||||||
|
After reasoning, if you find you lack some knowledge, you can call a search engine by <search> query </search>. \
|
||||||
|
Based on the user's core intent, formulate the most effective search query using specific, descriptive keywords that differentiate the topic clearly. \
|
||||||
|
Aim for queries that resemble how an expert searcher might phrase it, like using "compare lithium-ion vs solid-state battery efficiency" rather than just "batteries". \
|
||||||
|
The document will be provided inside <information> and </information> tags to you later. \
|
||||||
|
You can search as many turns as you want, but only one search query per turn. \
|
||||||
|
If you find no further external knowledge needed, you can directly provide the answer inside <answer> and </answer>, without detailed illustrations. \
|
||||||
|
Only answer when you have 100% confidence in the search results, else continue searching. \
|
||||||
|
Question: {q}\n"""
|
||||||
|
return user_prompt
|
||||||
|
|
||||||
|
|
||||||
|
def format_search_results(results: str | list[str]) -> str:
|
||||||
|
"""
|
||||||
|
Format search results for display, matching the format from infer.py.
|
||||||
|
Each result should be in the format: "Doc X(Title: Y) content"
|
||||||
|
|
||||||
|
Args:
|
||||||
|
results: Search results as string or list of strings
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
Formatted search results with document titles
|
||||||
|
"""
|
||||||
|
if isinstance(results, list):
|
||||||
|
# If results are already in the correct format, just join them
|
||||||
|
if any("Doc" in r and "Title:" in r for r in results):
|
||||||
|
content = "\n".join(results)
|
||||||
|
else:
|
||||||
|
# If results are raw content, format them with default titles
|
||||||
|
content = "\n".join([f"Doc {i + 1}(Title: Document {i + 1})\n{r}" for i, r in enumerate(results)])
|
||||||
|
else:
|
||||||
|
# If single result is already formatted, use it as is
|
||||||
|
if "Doc" in results and "Title:" in results:
|
||||||
|
content = results
|
||||||
|
else:
|
||||||
|
# If single result is raw content, format it with default title
|
||||||
|
content = f"Doc 1(Title: Document 1)\n{results}"
|
||||||
|
|
||||||
|
return content
|
||||||
@ -0,0 +1,312 @@
|
|||||||
|
"""
|
||||||
|
Reward functions for RL training.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import re
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from src.config import logger
|
||||||
|
from src.evaluation import check_student_answers
|
||||||
|
|
||||||
|
|
||||||
|
def build_reward_correctness_fn(
|
||||||
|
vllm_generate_func,
|
||||||
|
tokenizer,
|
||||||
|
):
|
||||||
|
"""Build a reward function that checks correctness of student answers.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
vllm_generate_func: Function to generate answers using vLLM
|
||||||
|
tokenizer: Tokenizer for the model
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
A reward function that takes prompts and completions and returns correctness scores
|
||||||
|
"""
|
||||||
|
|
||||||
|
def reward_correctness(prompts: list, completions: list, **reward_kwargs) -> list:
|
||||||
|
"""Calculate reward based on correctness of student answers.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompts: List of input prompts
|
||||||
|
completions: List of model completions
|
||||||
|
**reward_kwargs: Additional arguments for reward calculation
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of correctness scores between 0 and 1
|
||||||
|
"""
|
||||||
|
teacher_answers = reward_kwargs["answer"]
|
||||||
|
student_answers = [completion["messages"][-1]["content"] for completion in completions]
|
||||||
|
|
||||||
|
# Log non-exact matches
|
||||||
|
for i, (student, teacher) in enumerate(zip(student_answers, teacher_answers)):
|
||||||
|
if student.strip().lower() != teacher.strip().lower():
|
||||||
|
logger.debug(f"Non-exact match at index {i}:\nStudent: {student}\nTeacher: {teacher}")
|
||||||
|
|
||||||
|
correct = check_student_answers(
|
||||||
|
prompts,
|
||||||
|
teacher_answers,
|
||||||
|
student_answers,
|
||||||
|
vllm_generate_func=vllm_generate_func,
|
||||||
|
tokenizer=tokenizer,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Log correctness metrics with length info
|
||||||
|
logger.info(f"Correctness metrics: {correct}")
|
||||||
|
logger.info(f"Average correctness: {np.mean(correct):.2f}")
|
||||||
|
logger.info(f"Standard deviation: {np.std(correct):.2f}")
|
||||||
|
|
||||||
|
# Log length metrics
|
||||||
|
student_lengths = [len(ans.strip()) for ans in student_answers]
|
||||||
|
teacher_lengths = [len(ans.strip()) for ans in teacher_answers]
|
||||||
|
logger.info(f"Student lengths: {student_lengths}")
|
||||||
|
logger.info(f"Teacher lengths: {teacher_lengths}")
|
||||||
|
logger.info(f"Average student length: {np.mean(student_lengths):.2f}")
|
||||||
|
logger.info(f"Average teacher length: {np.mean(teacher_lengths):.2f}")
|
||||||
|
logger.info(f"Length ratio: {np.mean(student_lengths) / np.mean(teacher_lengths):.2f}")
|
||||||
|
|
||||||
|
return correct
|
||||||
|
|
||||||
|
return reward_correctness
|
||||||
|
|
||||||
|
|
||||||
|
def reward_format(prompts: list, completions: list, **reward_kwargs) -> list:
|
||||||
|
"""Reward function that checks if the completion follows the required format with proper tags.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompts: List of input prompts
|
||||||
|
completions: List of completion dictionaries containing messages
|
||||||
|
**reward_kwargs: Additional reward parameters
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list: List of rewards (1.0 for valid format, 0.0 for invalid)
|
||||||
|
"""
|
||||||
|
# Regex patterns for each tag type - no markdown allowed
|
||||||
|
think_pattern = r"<think>[\s\S]*?</think>"
|
||||||
|
search_pattern = r"<search>[\s\S]*?</search>"
|
||||||
|
answer_pattern = r"<answer>[\s\S]*?</answer>"
|
||||||
|
|
||||||
|
# Information tag patterns - handle multiple variants
|
||||||
|
info_patterns = [
|
||||||
|
r"<information>[\s\S]*?</information>", # Standard
|
||||||
|
r"<info>[\s\S]*?</info>", # Shortened
|
||||||
|
r"<Info[\w]*>[\s\S]*?</Info[\w]*>", # Capitalized variants
|
||||||
|
r"<INFORMATION>[\s\S]*?</INFORMATION>", # Uppercase
|
||||||
|
r"<INFO>[\s\S]*?</INFO>", # Uppercase shortened
|
||||||
|
]
|
||||||
|
|
||||||
|
# Invalid patterns (bold/italic tags)
|
||||||
|
invalid_patterns = [
|
||||||
|
r"\*\*<\/?(?:think|search|answer|information|info)>\*\*", # Bold tags
|
||||||
|
r"\*<\/?(?:think|search|answer|information|info)>\*", # Italic tags
|
||||||
|
r"_<\/?(?:think|search|answer|information|info)>_", # Underscore italic
|
||||||
|
]
|
||||||
|
|
||||||
|
rewards = []
|
||||||
|
|
||||||
|
for completion in completions:
|
||||||
|
messages = completion.get("messages", [])
|
||||||
|
assistant_msgs = [msg["content"] for msg in messages if msg["role"] == "assistant"]
|
||||||
|
|
||||||
|
if not assistant_msgs:
|
||||||
|
rewards.append(0.0)
|
||||||
|
continue
|
||||||
|
|
||||||
|
content = assistant_msgs[-1] # Get the last assistant message
|
||||||
|
|
||||||
|
# Check for invalid markdown formatting
|
||||||
|
has_invalid_tags = any(re.search(pattern, content) for pattern in invalid_patterns)
|
||||||
|
if has_invalid_tags:
|
||||||
|
logger.debug("Found markdown-formatted tags in response")
|
||||||
|
rewards.append(0.0)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Check for any information tag variants (should not exist in assistant messages)
|
||||||
|
has_info_tags = False
|
||||||
|
for pattern in info_patterns:
|
||||||
|
info_matches = re.findall(pattern, content, re.IGNORECASE)
|
||||||
|
if info_matches:
|
||||||
|
logger.debug(f"Found {len(info_matches)} information tag(s) of type '{pattern}' in assistant message")
|
||||||
|
has_info_tags = True
|
||||||
|
break
|
||||||
|
|
||||||
|
if has_info_tags:
|
||||||
|
rewards.append(0.0)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Find all tag matches
|
||||||
|
think_matches = re.findall(think_pattern, content)
|
||||||
|
search_matches = re.findall(search_pattern, content)
|
||||||
|
answer_matches = re.findall(answer_pattern, content)
|
||||||
|
|
||||||
|
# Verify tag presence and count
|
||||||
|
has_think = len(think_matches) >= 1
|
||||||
|
has_answer = len(answer_matches) == 1 # Must have exactly one answer
|
||||||
|
has_search = len(search_matches) >= 1 # One or more search tags
|
||||||
|
|
||||||
|
# Check for search and answer in the same message (not allowed)
|
||||||
|
if has_search and has_answer:
|
||||||
|
logger.debug("Found both search and answer tags in the same message")
|
||||||
|
rewards.append(0.0)
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Award reward - must have think tag and either answer or search (but not both)
|
||||||
|
reward = 1.0 if has_think and (has_answer or has_search) else 0.0
|
||||||
|
rewards.append(reward)
|
||||||
|
|
||||||
|
# Log issues for debugging
|
||||||
|
if not reward:
|
||||||
|
logger.debug(f"Format issues - think: {has_think}, answer: {has_answer}, search: {has_search}")
|
||||||
|
if search_matches:
|
||||||
|
logger.debug(f"Number of search tags: {len(search_matches)}")
|
||||||
|
|
||||||
|
# Log overall metrics
|
||||||
|
logger.info(f"Format reward metrics - Mean: {np.mean(rewards):.3f}, Valid formats: {sum(rewards)}/{len(rewards)}")
|
||||||
|
|
||||||
|
return rewards
|
||||||
|
|
||||||
|
|
||||||
|
# TODO: Implement this reward function if the project survives
|
||||||
|
def reward_long_query(completions, **kwargs):
|
||||||
|
"""Reward function that checks if the query is long."""
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
def reward_retry(prompts: list, completions: list, **reward_kwargs) -> list:
|
||||||
|
"""
|
||||||
|
Reward function that encourages optimal retry behavior.
|
||||||
|
Rewards increase with more search attempts but caps at optimal_search_count.
|
||||||
|
Penalizes having multiple searches in a single message.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompts: List of input prompts
|
||||||
|
completions: List of completion dictionaries with messages
|
||||||
|
**reward_kwargs: Additional reward parameters (chunk_id, answer, etc.)
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
List of rewards for each completion, rounded to 3 decimal places
|
||||||
|
"""
|
||||||
|
rewards = []
|
||||||
|
search_queries = []
|
||||||
|
violations = []
|
||||||
|
|
||||||
|
# Config for retry rewards
|
||||||
|
optimal_search_count = 5 # Cap rewards at this many searches
|
||||||
|
base_reward = 0.2 # Base reward for having at least one search
|
||||||
|
increment = 0.15 # Reward increment per search attempt (0.2 + 5*0.15 = 0.95 max)
|
||||||
|
violation_penalty = 0.5 # Penalty for having multiple searches in one message
|
||||||
|
|
||||||
|
# Regex pattern for search tags
|
||||||
|
search_pattern = r"<search>[\s\S]*?</search>"
|
||||||
|
|
||||||
|
for completion in completions:
|
||||||
|
# Get assistant messages
|
||||||
|
assistant_messages = [msg["content"] for msg in completion["messages"] if msg["role"] == "assistant"]
|
||||||
|
|
||||||
|
# Count search tags in assistant messages
|
||||||
|
message_searches = []
|
||||||
|
for msg in assistant_messages:
|
||||||
|
# Find all search tags in each message
|
||||||
|
search_matches = re.findall(search_pattern, msg)
|
||||||
|
message_searches.append(len(search_matches))
|
||||||
|
|
||||||
|
# Record total search queries
|
||||||
|
total_searches = sum(message_searches)
|
||||||
|
search_queries.append(total_searches)
|
||||||
|
|
||||||
|
# Check for violations (more than one search query per message)
|
||||||
|
violation = any(count > 1 for count in message_searches)
|
||||||
|
violations.append(violation)
|
||||||
|
|
||||||
|
# Calculate reward
|
||||||
|
if total_searches == 0:
|
||||||
|
reward = 0.0 # No searches = no reward
|
||||||
|
else:
|
||||||
|
# Base reward for having at least one search
|
||||||
|
reward = base_reward
|
||||||
|
|
||||||
|
# Add incremental reward for each search up to optimal_search_count
|
||||||
|
search_bonus = min(total_searches, optimal_search_count) * increment
|
||||||
|
reward += search_bonus
|
||||||
|
|
||||||
|
# Cap reward at 1.0
|
||||||
|
reward = min(1.0, reward)
|
||||||
|
|
||||||
|
# Apply penalty if there's a violation
|
||||||
|
if violation:
|
||||||
|
reward *= 1 - violation_penalty
|
||||||
|
|
||||||
|
# Round to 3 decimal places to avoid floating point precision issues
|
||||||
|
reward = round(reward, 3)
|
||||||
|
|
||||||
|
rewards.append(reward)
|
||||||
|
|
||||||
|
# Log metrics with search distribution info
|
||||||
|
logger.info(f"Retry behavior rewards: {np.mean(rewards):.3f} ± {np.std(rewards):.3f}")
|
||||||
|
logger.info(f"Search tags per completion: {np.mean(search_queries):.2f} ± {np.std(search_queries):.2f}")
|
||||||
|
logger.info(f"Violations (>1 search per message): {sum(violations)}/{len(violations)}")
|
||||||
|
logger.info(f"Search counts distribution: {search_queries}")
|
||||||
|
|
||||||
|
return rewards
|
||||||
|
|
||||||
|
|
||||||
|
def reward_em_chunk(prompts: list, completions: list, **reward_kwargs) -> list:
|
||||||
|
"""Reward function that checks if model's search queries hit the correct chunk content.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
prompts: List of input prompts
|
||||||
|
completions: List of completion dictionaries with messages
|
||||||
|
**reward_kwargs: Additional reward parameters including:
|
||||||
|
- chunk_content: List of correct chunk contents to match against
|
||||||
|
- step: Optional step number for logging metrics
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
list: List of rewards (1.0 for exact match, 0.0 otherwise)
|
||||||
|
|
||||||
|
Raises:
|
||||||
|
ValueError: If chunk_content is not provided in reward_kwargs
|
||||||
|
"""
|
||||||
|
logger.debug(f"Calculating rewards for {len(prompts)} prompts")
|
||||||
|
|
||||||
|
# Get correct chunk contents from reward kwargs
|
||||||
|
correct_contents = reward_kwargs.get("chunk_content", [])
|
||||||
|
if not correct_contents:
|
||||||
|
logger.error("No chunk_content provided in reward_kwargs")
|
||||||
|
raise ValueError("chunk_content must be provided in reward_kwargs")
|
||||||
|
|
||||||
|
rewards = []
|
||||||
|
for i, (completion, correct_content) in enumerate(zip(completions, correct_contents)):
|
||||||
|
# Get all messages from ipython or user roles that start with <information>
|
||||||
|
search_results = [
|
||||||
|
msg["content"]
|
||||||
|
for msg in completion["messages"]
|
||||||
|
if msg["role"] in ("ipython", "user") and msg["content"].strip().startswith("<information>")
|
||||||
|
]
|
||||||
|
logger.debug(f"Found {len(search_results)} search results for prompt {i}")
|
||||||
|
|
||||||
|
# Log ground truth and searched chunks for debugging
|
||||||
|
logger.info(f"📝 Ground Truth Chunk: {correct_content}")
|
||||||
|
for j, result in enumerate(search_results):
|
||||||
|
logger.info(f"🔍 Searched Chunk {j + 1}: {result}")
|
||||||
|
|
||||||
|
# Check if any search hit the correct chunk content
|
||||||
|
found_correct_chunk = any(correct_content in result for result in search_results)
|
||||||
|
|
||||||
|
if not found_correct_chunk:
|
||||||
|
logger.warning(
|
||||||
|
f"Failed to find correct chunk for prompt {i}:\n"
|
||||||
|
f"Search results: {[r[:100] + '...' for r in search_results]}"
|
||||||
|
)
|
||||||
|
|
||||||
|
reward = 1.0 if found_correct_chunk else 0.0
|
||||||
|
rewards.append(reward)
|
||||||
|
logger.debug(f"Reward for prompt {i}: {reward}")
|
||||||
|
|
||||||
|
# Log summary metrics
|
||||||
|
logger.info("Chunk Query Rewards Summary:")
|
||||||
|
logger.info(f"Total prompts: {len(prompts)}")
|
||||||
|
logger.info(f"Correct matches: {sum(rewards)}")
|
||||||
|
logger.info(f"Average reward: {np.mean(rewards):.3f}")
|
||||||
|
logger.info(f"Reward std: {np.std(rewards):.3f}")
|
||||||
|
|
||||||
|
return rewards
|
||||||
Loading…
Reference in new issue