File size: 4,451 Bytes
66407c5 |
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 |
# Copyright 2024 Bytedance Ltd. and/or its affiliates
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Create a simple multi-turn dataset for testing
"""
import argparse
import os
import pandas as pd
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--local_dir", default="~/data/multiturn")
parser.add_argument("--hdfs_dir", default=None)
args = parser.parse_args()
# Create example conversations
conversations = []
# Conversation 1
conversations.append(
{
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is the capital of France?"},
{"role": "assistant", "content": "The capital of France is Paris."},
{"role": "user", "content": "And what about Germany?"},
{"role": "assistant", "content": "The capital of Germany is Berlin."},
]
}
)
# Conversation 2
conversations.append(
{
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Can you explain quantum computing?"},
{
"role": "assistant",
"content": "Quantum computing is a type of computing that uses quantum-mechanical phenomena, such as superposition and entanglement, to perform operations on data.",
},
{"role": "user", "content": "How is it different from classical computing?"},
{
"role": "assistant",
"content": "Classical computing uses bits that are either 0 or 1, while quantum computing uses quantum bits or qubits that can exist in multiple states simultaneously due to superposition.",
},
]
}
)
# Conversation 3
conversations.append(
{
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Write a simple Python function to calculate factorial."},
{
"role": "assistant",
"content": "```python\ndef factorial(n):\n if n == 0 or n == 1:\n return 1\n else:\n return n * factorial(n-1)\n```\n\nThis is a recursive function to calculate the factorial of a number.",
},
{"role": "user", "content": "Can you make it iterative instead?"},
{
"role": "assistant",
"content": "```python\ndef factorial(n):\n result = 1\n for i in range(1, n+1):\n result *= i\n return result\n```\n\nThis is an iterative version of the factorial function.",
},
]
}
)
# Create train and test datasets
train_data = conversations[:2] # First 2 conversations for training
test_data = conversations[2:] # Last conversation for testing
# Create output directory
local_dir = os.path.expanduser(args.local_dir)
os.makedirs(local_dir, exist_ok=True)
# Save to parquet files
train_df = pd.DataFrame(train_data)
test_df = pd.DataFrame(test_data)
train_df.to_parquet(os.path.join(local_dir, "train.parquet"))
test_df.to_parquet(os.path.join(local_dir, "test.parquet"))
# Handle HDFS if specified
if args.hdfs_dir is not None:
try:
from verl.utils.hdfs_io import copy, makedirs
makedirs(args.hdfs_dir)
copy(src=local_dir, dst=args.hdfs_dir)
except ImportError:
print("Warning: HDFS support not available. Skipping HDFS copy.")
# Print statistics
print(f"Train dataset size: {len(train_df)}")
print(f"Test dataset size: {len(test_df)}")
print(f"Data saved to {local_dir}")
if __name__ == "__main__":
main()
|