|
|
"""
|
|
|
Tokenizer adapter implementations for different models.
|
|
|
This module provides adapter classes for handling different tokenizer formats.
|
|
|
"""
|
|
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
|
|
import torch
|
|
|
|
|
|
from config import logger
|
|
|
|
|
|
|
|
|
class TokenizerAdapter(ABC):
|
|
|
"""Base class for tokenizer adapters."""
|
|
|
|
|
|
@abstractmethod
|
|
|
def get_assistant_marker(self) -> str:
|
|
|
"""Get the assistant marker for the model."""
|
|
|
pass
|
|
|
|
|
|
@abstractmethod
|
|
|
def get_end_marker(self) -> str:
|
|
|
"""Get the end marker for the model."""
|
|
|
pass
|
|
|
|
|
|
@abstractmethod
|
|
|
def get_mask(self, text: str, tokenizer) -> torch.Tensor:
|
|
|
"""Get the mask for the model's response."""
|
|
|
pass
|
|
|
|
|
|
@abstractmethod
|
|
|
def split_prompt_assistant(self, text: str) -> tuple[str, str]:
|
|
|
"""Split conversation text into prompt and assistant response."""
|
|
|
pass
|
|
|
|
|
|
|
|
|
class LlamaTokenizerAdapter(TokenizerAdapter):
|
|
|
"""Adapter for Llama model tokenizer."""
|
|
|
|
|
|
def get_assistant_marker(self) -> str:
|
|
|
"""Get the assistant marker."""
|
|
|
return "<|start_header_id|>assistant<|end_header_id|>"
|
|
|
|
|
|
def get_end_marker(self) -> str:
|
|
|
"""Get the end marker."""
|
|
|
return "<|eot_id|>"
|
|
|
|
|
|
def split_prompt_assistant(self, convo_text: str) -> tuple[str, str]:
|
|
|
"""Split the text into prompt and assistant parts.
|
|
|
|
|
|
Args:
|
|
|
convo_text: The text to split
|
|
|
|
|
|
Returns:
|
|
|
A tuple of (prompt, assistant)
|
|
|
"""
|
|
|
# EXACT replication from rl_helpers.py but using existing method
|
|
|
marker = self.get_assistant_marker() # Use existing method but same value
|
|
|
idx = convo_text.find(marker)
|
|
|
if idx == -1:
|
|
|
raise ValueError("Could not find assistant marker in conversation text.")
|
|
|
return convo_text, ""
|
|
|
|
|
|
# Include the marker in the prompt by slicing up to the end of the marker.
|
|
|
prompt = convo_text[: idx + len(marker)]
|
|
|
# The assistant response is everything after the marker.
|
|
|
assistant_response = convo_text[idx + len(marker) :]
|
|
|
return prompt, assistant_response
|
|
|
|
|
|
def get_mask(self, text: str, tokenizer) -> torch.Tensor:
|
|
|
"""Get the mask for the text.
|
|
|
|
|
|
Args:
|
|
|
text: The text to get the mask for
|
|
|
tokenizer: The tokenizer to use
|
|
|
|
|
|
Returns:
|
|
|
A tensor of 0s and 1s where 1s indicate assistant tokens
|
|
|
"""
|
|
|
# Log input
|
|
|
logger.debug(f"🔍 Llama: Full text length: {len(text)}")
|
|
|
|
|
|
# EXACT replication from rl_helpers.py but using existing methods
|
|
|
encoding = tokenizer(text, add_special_tokens=False)
|
|
|
# Use existing methods but same values
|
|
|
start_header_id = tokenizer.convert_tokens_to_ids("<|start_header_id|>")
|
|
|
assistant_token = tokenizer.convert_tokens_to_ids("assistant")
|
|
|
end_header_id = tokenizer.convert_tokens_to_ids("<|end_header_id|>")
|
|
|
eot_id = tokenizer.convert_tokens_to_ids(self.get_end_marker()) # Use existing method but same value
|
|
|
|
|
|
# Log token IDs
|
|
|
logger.debug(f"🔍 Llama: Tokenized length: {len(encoding.input_ids)}")
|
|
|
logger.debug(f"🔍 Llama: Input IDs: {encoding.input_ids}")
|
|
|
logger.debug(
|
|
|
f"🔍 Llama: Special token IDs: start={start_header_id}, assistant={assistant_token}, end={end_header_id}, eot={eot_id}"
|
|
|
)
|
|
|
|
|
|
assistant_ranges = []
|
|
|
i = 0
|
|
|
while i < len(encoding.input_ids) - 1:
|
|
|
if encoding.input_ids[i] == start_header_id and encoding.input_ids[i + 1] == assistant_token:
|
|
|
logger.debug(f"🔍 Llama: Found assistant marker at position {i}")
|
|
|
logger.debug(f"🔍 Llama: Assistant marker tokens: {encoding.input_ids[i : i + 2]}")
|
|
|
i += 2
|
|
|
while i < len(encoding.input_ids) and encoding.input_ids[i] != end_header_id:
|
|
|
i += 1
|
|
|
i += 2
|
|
|
start_idx = i
|
|
|
logger.debug(f"🔍 Llama: Found start of response at {start_idx}")
|
|
|
logger.debug(f"🔍 Llama: Start token ID: {encoding.input_ids[start_idx]}")
|
|
|
while i < len(encoding.input_ids) and encoding.input_ids[i] != eot_id:
|
|
|
i += 1
|
|
|
end_idx = i
|
|
|
logger.debug(f"🔍 Llama: Found end of response at {end_idx}")
|
|
|
logger.debug(f"🔍 Llama: End token ID: {encoding.input_ids[end_idx]}")
|
|
|
logger.debug(f"🔍 Llama: Response token IDs: {encoding.input_ids[start_idx:end_idx]}")
|
|
|
assistant_ranges.append((start_idx, end_idx))
|
|
|
else:
|
|
|
i += 1
|
|
|
|
|
|
mask = [0] * len(encoding.input_ids)
|
|
|
for start_idx, end_idx in assistant_ranges:
|
|
|
for idx in range(start_idx, end_idx):
|
|
|
mask[idx] = 1
|
|
|
|
|
|
mask = torch.tensor(mask, dtype=torch.int)
|
|
|
|
|
|
# Log final mask
|
|
|
logger.debug(f"🔍 Llama: Final mask shape: {mask.shape}")
|
|
|
logger.debug(f"🔍 Llama: Mask sum: {mask.sum().item()}")
|
|
|
logger.debug(f"🔍 Llama: Mask: {mask}")
|
|
|
|
|
|
# Additional debug info
|
|
|
try:
|
|
|
prompt, response = self.split_prompt_assistant(text)
|
|
|
prompt_tokens = tokenizer(prompt, add_special_tokens=False).input_ids
|
|
|
response_tokens = tokenizer(response, add_special_tokens=False).input_ids
|
|
|
|
|
|
logger.debug(f"🔍 Llama: Prompt length: {len(prompt)}")
|
|
|
logger.debug(f"🔍 Llama: Response length: {len(response)}")
|
|
|
logger.debug(f"🔍 Llama: Prompt token IDs: {prompt_tokens}")
|
|
|
logger.debug(f"🔍 Llama: Response token IDs: {response_tokens}")
|
|
|
logger.debug(f"🔍 Llama: Prompt: {prompt[:100]}...")
|
|
|
logger.debug(f"🔍 Llama: Response: {response[:100]}...")
|
|
|
logger.debug(f"🔍 Llama: Full input IDs length: {len(encoding.input_ids)}")
|
|
|
logger.debug(f"🔍 Llama: Prompt + Response token IDs length: {len(prompt_tokens) + len(response_tokens)}")
|
|
|
logger.debug(
|
|
|
f"🔍 Llama: Difference in lengths: {len(encoding.input_ids) - (len(prompt_tokens) + len(response_tokens))}"
|
|
|
)
|
|
|
except Exception as e:
|
|
|
logger.error(f"🔍 Llama: Error splitting prompt/response: {e}")
|
|
|
|
|
|
return mask
|
|
|
|
|
|
|
|
|
class R1DistilTokenizerAdapter(TokenizerAdapter):
|
|
|
"""Adapter for R1-Distil model tokenizer."""
|
|
|
|
|
|
def get_assistant_marker(self) -> str:
|
|
|
marker = "<｜Assistant｜>"
|
|
|
return marker
|
|
|
|
|
|
def get_end_marker(self) -> str:
|
|
|
marker = "<｜end▁of▁sentence｜>"
|
|
|
return marker
|
|
|
|
|
|
def get_begin_marker(self) -> str:
|
|
|
return "<｜begin▁of▁sentence｜>"
|
|
|
|
|
|
def get_user_marker(self) -> str:
|
|
|
return "<｜User｜>"
|
|
|
|
|
|
def get_mask(self, text: str, tokenizer) -> torch.Tensor:
|
|
|
"""Get the mask for the text.
|
|
|
|
|
|
Args:
|
|
|
text: The text to get the mask for
|
|
|
tokenizer: The tokenizer to use
|
|
|
|
|
|
Returns:
|
|
|
A tensor of 0s and 1s where 1s indicate assistant tokens
|
|
|
"""
|
|
|
logger.debug(f"🔍 R1Distil: Getting mask for text length: {len(text)}")
|
|
|
|
|
|
# Get all markers
|
|
|
assistant_marker = self.get_assistant_marker()
|
|
|
end_marker = self.get_end_marker()
|
|
|
|
|
|
# Get the full tokenization
|
|
|
encoding = tokenizer(text, add_special_tokens=False)
|
|
|
tokens = encoding.input_ids
|
|
|
logger.debug(f"🔍 R1Distil: Full text token IDs: {tokens}")
|
|
|
|
|
|
# Create mask initialized to zeros - ENSURE SAME LENGTH AS INPUT_IDS
|
|
|
mask = torch.zeros(len(tokens), dtype=torch.int)
|
|
|
|
|
|
# Get token IDs for markers
|
|
|
assistant_tokens = tokenizer(assistant_marker, add_special_tokens=False).input_ids
|
|
|
end_tokens = tokenizer(end_marker, add_special_tokens=False).input_ids
|
|
|
logger.debug(f"🔍 R1Distil: Assistant marker token IDs: {assistant_tokens}")
|
|
|
logger.debug(f"🔍 R1Distil: End marker token IDs: {end_tokens}")
|
|
|
|
|
|
# Find all assistant responses
|
|
|
assistant_ranges = []
|
|
|
i = 0
|
|
|
while i < len(tokens):
|
|
|
# Look for assistant marker
|
|
|
if i + len(assistant_tokens) <= len(tokens) and tokens[i : i + len(assistant_tokens)] == assistant_tokens:
|
|
|
logger.debug(f"🔍 R1Distil: Found assistant marker at position {i}")
|
|
|
|
|
|
# Start masking AFTER the assistant marker
|
|
|
start_idx = i + len(assistant_tokens)
|
|
|
|
|
|
# Find end marker
|
|
|
end_idx = None
|
|
|
j = start_idx
|
|
|
while j < len(tokens):
|
|
|
if j + len(end_tokens) <= len(tokens) and tokens[j : j + len(end_tokens)] == end_tokens:
|
|
|
end_idx = j # Don't include the end marker in the mask
|
|
|
break
|
|
|
j += 1
|
|
|
|
|
|
if end_idx is None:
|
|
|
# If no end marker found, mask until the end
|
|
|
end_idx = len(tokens)
|
|
|
|
|
|
logger.debug(f"🔍 R1Distil: Response range: {start_idx} to {end_idx}")
|
|
|
assistant_ranges.append((start_idx, end_idx))
|
|
|
i = end_idx + len(end_tokens) # Start next search after the end marker
|
|
|
else:
|
|
|
i += 1
|
|
|
|
|
|
# Apply mask for all found ranges
|
|
|
for start_idx, end_idx in assistant_ranges:
|
|
|
mask[start_idx:end_idx] = 1
|
|
|
|
|
|
logger.debug(f"🔍 R1Distil: Found {len(assistant_ranges)} assistant responses")
|
|
|
logger.debug(f"🔍 R1Distil: Final mask sum: {mask.sum().item()}")
|
|
|
logger.debug(f"🔍 R1Distil: Final mask length: {len(mask)}")
|
|
|
logger.debug(f"🔍 R1Distil: Mask: {mask}")
|
|
|
|
|
|
return mask
|
|
|
|
|
|
def split_prompt_assistant(self, text: str) -> tuple[str, str]:
|
|
|
"""Split the text into prompt and assistant parts.
|
|
|
|
|
|
Args:
|
|
|
text: The text to split
|
|
|
|
|
|
Returns:
|
|
|
A tuple of (prompt, assistant)
|
|
|
"""
|
|
|
logger.debug(f"🔍 R1Distil: Splitting text of length: {len(text)}")
|
|
|
|
|
|
# Find the assistant marker
|
|
|
marker = self.get_assistant_marker()
|
|
|
end_marker = self.get_end_marker()
|
|
|
|
|
|
# Find ALL assistant markers in the text
|
|
|
assistant_markers = []
|
|
|
pos = 0
|
|
|
while True:
|
|
|
pos = text.find(marker, pos)
|
|
|
if pos == -1:
|
|
|
break
|
|
|
assistant_markers.append(pos)
|
|
|
pos += len(marker)
|
|
|
|
|
|
if not assistant_markers:
|
|
|
raise ValueError("Could not find assistant marker in text")
|
|
|
|
|
|
# Get the positions of all markers for later use
|
|
|
marker_positions = []
|
|
|
for start_pos in assistant_markers:
|
|
|
response_start = start_pos + len(marker)
|
|
|
|
|
|
# Find the end marker after this response
|
|
|
end_pos = text.find(end_marker, response_start)
|
|
|
if end_pos == -1:
|
|
|
end_pos = len(text)
|
|
|
else:
|
|
|
end_pos = end_pos + len(end_marker)
|
|
|
|
|
|
marker_positions.append((start_pos, response_start, end_pos))
|
|
|
|
|
|
# Get the full response (all assistant outputs concatenated)
|
|
|
full_response = ""
|
|
|
for _, resp_start, resp_end in marker_positions:
|
|
|
full_response += text[resp_start:resp_end]
|
|
|
|
|
|
# Include ALL assistant markers and responses in the response
|
|
|
# This matches how the mask is generated in get_mask
|
|
|
first_assistant_pos = marker_positions[0][0]
|
|
|
last_response_end = marker_positions[-1][2]
|
|
|
|
|
|
# Split into prompt and response
|
|
|
prompt = text[:first_assistant_pos] # Everything before the first assistant marker
|
|
|
response = text[first_assistant_pos:last_response_end] # All markers and responses
|
|
|
|
|
|
logger.debug(f"🔍 R1Distil: Prompt length: {len(prompt)}")
|
|
|
logger.debug(f"🔍 R1Distil: Response length: {len(response)}")
|
|
|
logger.debug(f"🔍 R1Distil: Response token count estimate: {len(response) / 4}") # Rough estimate
|
|
|
logger.debug(f"🔍 R1Distil: Final prompt: {prompt[:100]}...")
|
|
|
logger.debug(f"🔍 R1Distil: Final response: {response[:100]}...")
|
|
|
|
|
|
return prompt, response
|
|
|
|
|
|
|
|
|
class QwenTokenizerAdapter(TokenizerAdapter):
|
|
|
"""Adapter for Qwen2.5 model tokenizer."""
|
|
|
|
|
|
def get_assistant_marker(self) -> str:
|
|
|
"""Get the assistant marker."""
|
|
|
return "<|im_start|>assistant"
|
|
|
|
|
|
def get_end_marker(self) -> str:
|
|
|
"""Get the end marker."""
|
|
|
return "<|im_end|>"
|
|
|
|
|
|
def split_prompt_assistant(self, convo_text: str) -> tuple[str, str]:
|
|
|
"""Split the text into prompt and assistant parts.
|
|
|
|
|
|
Args:
|
|
|
convo_text: The text to split
|
|
|
|
|
|
Returns:
|
|
|
A tuple of (prompt, assistant)
|
|
|
"""
|
|
|
marker = self.get_assistant_marker()
|
|
|
idx = convo_text.find(marker)
|
|
|
if idx == -1:
|
|
|
raise ValueError("Could not find assistant marker in conversation text.")
|
|
|
return convo_text, ""
|
|
|
|
|
|
# Include the marker in the prompt by slicing up to the end of the marker
|
|
|
prompt = convo_text[: idx + len(marker)]
|
|
|
# The assistant response is everything after the marker
|
|
|
assistant_response = convo_text[idx + len(marker) :]
|
|
|
return prompt, assistant_response
|
|
|
|
|
|
def get_mask(self, text: str, tokenizer) -> torch.Tensor:
|
|
|
"""Get the mask for the text.
|
|
|
|
|
|
Args:
|
|
|
text: The text to get the mask for
|
|
|
tokenizer: The tokenizer to use
|
|
|
|
|
|
Returns:
|
|
|
A tensor of 0s and 1s where 1s indicate assistant tokens
|
|
|
"""
|
|
|
# Log input
|
|
|
logger.debug(f"🔍 Qwen: Full text length: {len(text)}")
|
|
|
|
|
|
encoding = tokenizer(text, add_special_tokens=False)
|
|
|
# Get token IDs for markers
|
|
|
im_start = tokenizer.convert_tokens_to_ids("<|im_start|>")
|
|
|
assistant_token = tokenizer.convert_tokens_to_ids("assistant")
|
|
|
im_end = tokenizer.convert_tokens_to_ids(self.get_end_marker())
|
|
|
|
|
|
# Log token IDs
|
|
|
logger.debug(f"🔍 Qwen: Tokenized length: {len(encoding.input_ids)}")
|
|
|
logger.debug(f"🔍 Qwen: Input IDs: {encoding.input_ids}")
|
|
|
logger.debug(f"🔍 Qwen: Special token IDs: im_start={im_start}, assistant={assistant_token}, im_end={im_end}")
|
|
|
|
|
|
assistant_ranges = []
|
|
|
i = 0
|
|
|
while i < len(encoding.input_ids) - 1:
|
|
|
if encoding.input_ids[i] == im_start and encoding.input_ids[i + 1] == assistant_token:
|
|
|
logger.debug(f"🔍 Qwen: Found assistant marker at position {i}")
|
|
|
logger.debug(f"🔍 Qwen: Assistant marker tokens: {encoding.input_ids[i : i + 2]}")
|
|
|
i += 2 # Skip past <|im_start|>assistant
|
|
|
start_idx = i
|
|
|
logger.debug(f"🔍 Qwen: Found start of response at {start_idx}")
|
|
|
logger.debug(f"🔍 Qwen: Start token ID: {encoding.input_ids[start_idx]}")
|
|
|
|
|
|
while i < len(encoding.input_ids) and encoding.input_ids[i] != im_end:
|
|
|
i += 1
|
|
|
end_idx = i
|
|
|
logger.debug(f"🔍 Qwen: Found end of response at {end_idx}")
|
|
|
logger.debug(f"🔍 Qwen: End token ID: {encoding.input_ids[end_idx]}")
|
|
|
logger.debug(f"🔍 Qwen: Response token IDs: {encoding.input_ids[start_idx:end_idx]}")
|
|
|
assistant_ranges.append((start_idx, end_idx))
|
|
|
else:
|
|
|
i += 1
|
|
|
|
|
|
mask = [0] * len(encoding.input_ids)
|
|
|
for start_idx, end_idx in assistant_ranges:
|
|
|
for idx in range(start_idx, end_idx):
|
|
|
mask[idx] = 1
|
|
|
|
|
|
mask = torch.tensor(mask, dtype=torch.int)
|
|
|
|
|
|
# Log final mask
|
|
|
logger.debug(f"🔍 Qwen: Final mask shape: {mask.shape}")
|
|
|
logger.debug(f"🔍 Qwen: Mask sum: {mask.sum().item()}")
|
|
|
logger.debug(f"🔍 Qwen: Mask: {mask}")
|
|
|
|
|
|
# Additional debug info
|
|
|
try:
|
|
|
prompt, response = self.split_prompt_assistant(text)
|
|
|
prompt_tokens = tokenizer(prompt, add_special_tokens=False).input_ids
|
|
|
response_tokens = tokenizer(response, add_special_tokens=False).input_ids
|
|
|
|
|
|
logger.debug(f"🔍 Qwen: Prompt length: {len(prompt)}")
|
|
|
logger.debug(f"🔍 Qwen: Response length: {len(response)}")
|
|
|
logger.debug(f"🔍 Qwen: Prompt token IDs: {prompt_tokens}")
|
|
|
logger.debug(f"🔍 Qwen: Response token IDs: {response_tokens}")
|
|
|
logger.debug(f"🔍 Qwen: Prompt: {prompt[:100]}...")
|
|
|
logger.debug(f"🔍 Qwen: Response: {response[:100]}...")
|
|
|
logger.debug(f"🔍 Qwen: Full input IDs length: {len(encoding.input_ids)}")
|
|
|
logger.debug(f"🔍 Qwen: Prompt + Response token IDs length: {len(prompt_tokens) + len(response_tokens)}")
|
|
|
logger.debug(
|
|
|
f"🔍 Qwen: Difference in lengths: {len(encoding.input_ids) - (len(prompt_tokens) + len(response_tokens))}"
|
|
|
)
|
|
|
except Exception as e:
|
|
|
logger.error(f"🔍 Qwen: Error splitting prompt/response: {e}")
|
|
|
|
|
|
return mask
|
