mirror of
https://github.com/huggingface/trl.git
synced 2025-10-21 02:53:59 +08:00
235 lines
9.5 KiB
Python
235 lines
9.5 KiB
Python
# Copyright 2020-2025 The HuggingFace Team. All rights reserved.
|
||
#
|
||
# 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.
|
||
|
||
from dataclasses import dataclass, field
|
||
|
||
from datasets import Dataset
|
||
from transformers import HfArgumentParser
|
||
from transformers.utils import get_json_schema
|
||
|
||
|
||
@dataclass
|
||
class ScriptArguments:
|
||
r"""
|
||
Arguments for the script.
|
||
|
||
Args:
|
||
test_size (`float`, *optional*, defaults to `0.1`):
|
||
Fraction of the dataset to include in the test split.
|
||
push_to_hub (`bool`, *optional*, defaults to `False`):
|
||
Whether to push the dataset to the Hugging Face Hub.
|
||
repo_id (`str`, *optional*, defaults to `"trl-internal-testing/zen"`):
|
||
Hugging Face repository ID to push the dataset to.
|
||
"""
|
||
|
||
test_size: float = field(
|
||
default=0.1,
|
||
metadata={"help": "Fraction of the dataset to include in the test split."},
|
||
)
|
||
push_to_hub: bool = field(
|
||
default=False,
|
||
metadata={"help": "Whether to push the dataset to the Hugging Face Hub."},
|
||
)
|
||
repo_id: str = field(
|
||
default="trl-internal-testing/toolcall",
|
||
metadata={"help": "Hugging Face repository ID to push the dataset to."},
|
||
)
|
||
|
||
|
||
def main(test_size, push_to_hub, repo_id):
|
||
# Fictitious functions to simulate tool calls
|
||
def start_timer(duration: int) -> int:
|
||
"""
|
||
Starts a timer for the specified duration in seconds.
|
||
|
||
Args:
|
||
duration: Duration in seconds to set the timer for.
|
||
|
||
Returns:
|
||
The duration set for the timer.
|
||
"""
|
||
return duration
|
||
|
||
def get_current_time(location: str) -> str:
|
||
"""
|
||
Returns the current time in the specified location.
|
||
|
||
Args:
|
||
location: The location for which to get the current time.
|
||
|
||
Returns:
|
||
The current time in the specified location.
|
||
"""
|
||
return "06:22:48"
|
||
|
||
def get_air_quality_index(location: str) -> int:
|
||
"""
|
||
Returns the air quality index for the specified location.
|
||
|
||
Args:
|
||
location: The location for which to get the air quality index.
|
||
|
||
Returns:
|
||
The air quality index for the specified location.
|
||
"""
|
||
return 53
|
||
|
||
def play_music(title: str, artist: str) -> dict:
|
||
"""
|
||
Plays music by the specified title and artist.
|
||
|
||
Args:
|
||
title: The title of the music to play.
|
||
artist: The artist of the music to play.
|
||
|
||
Returns:
|
||
A dictionary indicating the status of the music playback.
|
||
"""
|
||
return {"status": "Playing"}
|
||
|
||
def get_weather_forecast(city: str, date: str) -> dict:
|
||
"""
|
||
Returns the weather forecast for the specified city and date.
|
||
|
||
Args:
|
||
city: The city for which to get the weather forecast.
|
||
date: The date for which to get the weather forecast.
|
||
|
||
Returns:
|
||
A dictionary containing the temperature and weather condition.
|
||
"""
|
||
return {"temperature": 22, "condition": "partly cloudy"}
|
||
|
||
def control_light(room: str, state: str) -> dict:
|
||
"""
|
||
Controls the light in the specified room.
|
||
|
||
Args:
|
||
room: The room where the light should be controlled.
|
||
state: The desired state of the light ("on" or "off").
|
||
|
||
Returns:
|
||
A dictionary indicating the state of the light.
|
||
"""
|
||
return {"state": state}
|
||
|
||
def create_reminder(time: str, note: str) -> str:
|
||
"""
|
||
Creates a reminder for the specified time and note.
|
||
|
||
Args:
|
||
time: The time for the reminder.
|
||
note: The note for the reminder.
|
||
|
||
Returns:
|
||
A confirmation message indicating that the reminder has been set.
|
||
"""
|
||
return "I'll remind you to call mom at 7 PM."
|
||
|
||
def get_wind_conditions(city: str, unit: str) -> tuple[int, str]:
|
||
"""
|
||
Returns the wind conditions for the specified city.
|
||
|
||
Args:
|
||
city: The city for which to get the wind conditions.
|
||
unit: The unit of measurement for the wind speed (e.g., "mph").
|
||
|
||
Returns:
|
||
A tuple containing the wind speed and direction.
|
||
"""
|
||
return 14, "NW"
|
||
|
||
start_timer = get_json_schema(start_timer)
|
||
get_current_time = get_json_schema(get_current_time)
|
||
get_air_quality_index = get_json_schema(get_air_quality_index)
|
||
play_music = get_json_schema(play_music)
|
||
get_weather_forecast = get_json_schema(get_weather_forecast)
|
||
control_light = get_json_schema(control_light)
|
||
create_reminder = get_json_schema(create_reminder)
|
||
get_wind_conditions = get_json_schema(get_wind_conditions)
|
||
|
||
# fmt: off
|
||
dataset = Dataset.from_dict({
|
||
"messages": [
|
||
[
|
||
{"role": "user", "content": "Set a timer for 10 minutes."},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "start_timer", "arguments": {"duration": 600}}}]},
|
||
{"role": "tool", "name": "start_timer", "content": "600"},
|
||
{"role": "assistant", "content": "Timer set for 10 minutes."},
|
||
],
|
||
[
|
||
{"role": "user", "content": "What time is it in Tokyo?"},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "get_current_time", "arguments": {"location": "Tokyo"}}}]},
|
||
{"role": "tool", "name": "get_current_time", "content": "06:22:48"},
|
||
{"role": "assistant", "content": "The current time in Tokyo is 06:22 AM."},
|
||
],
|
||
[
|
||
{"role": "user", "content": "Is the air clean today in Lisbon?"},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "get_air_quality", "arguments": {"location": "Lisbon, Portugal"}}}]},
|
||
{"role": "tool", "name": "get_air_quality", "content": "53"},
|
||
{"role": "assistant", "content": "The air quality is moderate."},
|
||
],
|
||
[
|
||
{"role": "user", "content": "Play some music."},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "play_music", "arguments": {"title": "Take Five", "artist": "Dave Brubeck"}}}]},
|
||
{"role": "tool", "name": "play_music", "content": "{'status': 'Playing'}"},
|
||
{"role": "assistant", "content": "Enjoy the jazz tunes!"},
|
||
],
|
||
[
|
||
{"role": "user", "content": "What's the weather like tomorrow in Berlin?"},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "get_weather_forecast", "arguments": {"city": "Berlin", "date": "2025-06-16"}}}]},
|
||
{"role": "tool", "name": "get_weather_forecast", "content": "{'temperature': 22, 'condition': 'partly cloudy'}"},
|
||
{"role": "assistant", "content": "Tomorrow in Berlin will be partly cloudy with a high of 22°C."}
|
||
],
|
||
[
|
||
{"role": "user", "content": "Turn on the living room lights."},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "control_light", "arguments": {"room": "living room", "state": "on"}}}]},
|
||
{"role": "tool", "name": "control_light", "content": "{'state': 'on'}"},
|
||
{"role": "assistant", "content": "The living room lights are now on."}
|
||
],
|
||
[
|
||
{"role": "user", "content": "Remind me to call mom at 7 PM."},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "create_reminder", "arguments": {"time": "19:00", "note": "Call mom"}}}]},
|
||
{"role": "tool", "name": "create_reminder", "content": "Reminder set"},
|
||
{"role": "assistant", "content": "Okay, I’ll remind you to call mom at 7 PM."}
|
||
],
|
||
[
|
||
{"role": "user", "content": "How strong is the wind in Chicago right now?"},
|
||
{"role": "assistant", "tool_calls": [{"type": "function", "function": {"name": "get_wind_conditions", "arguments": {"city": "Chicago", "unit": "mph"}}}]},
|
||
{"role": "tool", "name": "get_wind_conditions", "content": "(14, 'NW')"},
|
||
{"role": "assistant", "content": "The wind in Chicago is blowing at 14 mph from the northwest."}
|
||
]
|
||
],
|
||
"tools": [
|
||
[start_timer, create_reminder],
|
||
[get_current_time],
|
||
[get_air_quality_index, get_weather_forecast, get_wind_conditions],
|
||
[play_music, control_light],
|
||
[get_weather_forecast, get_wind_conditions],
|
||
[control_light],
|
||
[start_timer, create_reminder],
|
||
[get_weather_forecast, get_wind_conditions],
|
||
]
|
||
})
|
||
dataset = dataset.train_test_split(test_size=test_size, shuffle=False)
|
||
if push_to_hub:
|
||
dataset.push_to_hub(repo_id)
|
||
# fmt: on
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = HfArgumentParser(ScriptArguments)
|
||
script_args = parser.parse_args_into_dataclasses()[0]
|
||
main(script_args.test_size, script_args.push_to_hub, script_args.repo_id)
|