File size: 6,380 Bytes
6f09d40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
#!/usr/bin/env python3
"""
Data preprocessing script.
Convert the generated dataset into a format directly consumable by SFTTrainer.
FunctionGemma expects a specific chat template structure.
Usage:
python -m src.prepare_dataset --input ./data/training_data.json --output ./data/prepared_dataset.json
"""
import json
import argparse
from pathlib import Path
from typing import List, Dict, Any
PROJECT_ROOT = Path(__file__).resolve().parent.parent
DEFAULT_INPUT = PROJECT_ROOT / "data" / "training_data.json"
DEFAULT_OUTPUT = PROJECT_ROOT / "data" / "prepared_dataset.json"
def convert_tool_calls_to_text(tool_calls: List[Dict]) -> str:
"""Convert tool_calls into plain text (FunctionGemma format)."""
if not tool_calls:
return ""
result_parts = []
for tc in tool_calls:
func = tc.get("function", {})
name = func.get("name", "")
args = func.get("arguments", {})
# FunctionGemma format: functionName(arguments)
args_str = json.dumps(args, ensure_ascii=False)
result_parts.append(f"{name}({args_str})")
return "\n".join(result_parts)
def convert_messages_for_sft(messages: List[Dict], tools: List[Dict] = None) -> List[Dict]:
"""
Convert message format for SFTTrainer.
Input:
[
{"role": "developer", "content": "..."},
{"role": "user", "content": "..."},
{"role": "assistant", "tool_calls": [...]} or {"role": "assistant", "content": "..."}
]
Output:
[
{"role": "system", "content": "..."}, # developer -> system
{"role": "user", "content": "..."},
{"role": "assistant", "content": "..."} # tool_calls flattened to text
]
"""
converted = []
# Build tools description
tools_description = ""
if tools:
tools_desc_parts = []
for tool in tools:
if tool.get("type") == "function":
func = tool.get("function", {})
name = func.get("name", "")
desc = func.get("description", "")
params = func.get("parameters", {})
tools_desc_parts.append(f"- {name}: {desc}")
if tools_desc_parts:
tools_description = "\n\nAvailable tools:\n" + "\n".join(tools_desc_parts)
for msg in messages:
role = msg.get("role", "")
if role == "developer":
# developer -> system
content = msg.get("content", "")
if tools_description:
content = content + tools_description
converted.append({
"role": "system",
"content": content
})
elif role == "user":
converted.append({
"role": "user",
"content": msg.get("content", "")
})
elif role == "assistant":
if "tool_calls" in msg:
# Convert tool_calls to text
tool_calls_text = convert_tool_calls_to_text(msg["tool_calls"])
converted.append({
"role": "assistant",
"content": tool_calls_text
})
else:
converted.append({
"role": "assistant",
"content": msg.get("content", "")
})
elif role == "tool":
# Tool response
converted.append({
"role": "tool",
"content": msg.get("content", "")
})
return converted
def prepare_dataset(input_path: str, output_path: str, format_type: str = "messages"):
"""
Prepare dataset.
format_type:
- "messages": output {"messages": [...]}
- "text": output {"text": "..."} (flattened text)
"""
print(f"Loading dataset: {input_path}")
with open(input_path, 'r', encoding='utf-8') as f:
data = json.load(f)
print(f"Raw samples: {len(data)}")
prepared_data = []
for i, item in enumerate(data):
messages = item.get("messages", [])
tools = item.get("tools", [])
# Convert messages
converted_messages = convert_messages_for_sft(messages, tools)
if format_type == "messages":
prepared_data.append({
"messages": converted_messages
})
elif format_type == "text":
# Convert to plain text
text_parts = []
for msg in converted_messages:
role = msg["role"]
content = msg["content"]
if role == "system":
text_parts.append(f"<start_of_turn>system\n{content}<end_of_turn>")
elif role == "user":
text_parts.append(f"<start_of_turn>user\n{content}<end_of_turn>")
elif role == "assistant":
text_parts.append(f"<start_of_turn>model\n{content}<end_of_turn>")
prepared_data.append({
"text": "\n".join(text_parts)
})
print(f"Processed samples: {len(prepared_data)}")
# Save
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(prepared_data, f, ensure_ascii=False, indent=2)
print(f"Saved to: {output_path}")
# Show example
print("\n" + "=" * 60)
print("Example:")
print("=" * 60)
if format_type == "messages":
example = prepared_data[0]
for msg in example["messages"]:
print(f"\n[{msg['role']}]")
print(msg["content"][:200] + "..." if len(msg["content"]) > 200 else msg["content"])
else:
print(prepared_data[0]["text"][:500] + "...")
return prepared_data
def main():
parser = argparse.ArgumentParser(description="Dataset preparation")
parser.add_argument("--input", type=str, default=str(DEFAULT_INPUT), help="Input file path")
parser.add_argument("--output", type=str, default=str(DEFAULT_OUTPUT), help="Output file path")
parser.add_argument("--format", type=str, choices=["messages", "text"], default="messages", help="Output format")
args = parser.parse_args()
prepare_dataset(args.input, args.output, args.format)
if __name__ == "__main__":
main()
|