Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add @file syntax to ds_tool #28

Merged
merged 3 commits into from
Jun 13, 2024
Merged

Add @file syntax to ds_tool #28

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4225172
Add @file syntax to ds_tool
juberti Jun 13, 2024
236a2ef
x
juberti Jun 13, 2024
3f3a6e1
cr
juberti Jun 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 41 additions & 52 deletions ultravox/tools/ds_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,18 @@ def __post_init__(self):
if self.audio_column_name is None:
self.audio_column_name = f"{self.column_name}_audio"

def map_split(self, ds_split: datasets.Dataset, num_proc: int) -> datasets.Dataset:
print(f'TTS mapping "{self.column_name}" to "{self.audio_column_name}"...')
return ds_split.map(self._map_sample, num_proc=num_proc).cast_column(
self.audio_column_name, datasets.Audio(sampling_rate=self.sample_rate)
)

def _map_sample(self, sample):
text = sample[self.column_name]
text = text["text"] if isinstance(text, dict) else text
sample[self.audio_column_name] = tts_client.tts(text)
return sample


@dataclasses.dataclass
class TextGenerationTask:
Expand All @@ -41,16 +53,35 @@ class TextGenerationTask:
max_tokens: int = 128
temperature: float = 0

def __post_init__(self):
if self.template.startswith("@"):
with open(self.template[1:], "r") as template_file:
self.template = template_file.read()

def map_split(self, ds_split: datasets.Dataset, num_proc: int) -> datasets.Dataset:
print(f'Generating "{self.new_column_name}" with template:\n{self.template}')
return ds_split.map(self._map_sample, num_proc=num_proc)

def _map_sample(self, sample):
input_text = self.template.format(**sample)
response = chat_client.chat.completions.create(
model=self.language_model,
messages=[{"role": "user", "content": input_text}],
max_tokens=self.max_tokens,
temperature=self.temperature,
)
sample[self.new_column_name] = response.choices[0].message.content
return sample


# This script is used to either generate audio samples from text using a TTS model, or to generate text samples using a text generation model.
# Ex: just ds_tool -t tts -d google/boolq -u fixie-ai/boolq-audio -c question -a audio --token $HF_WRITE_TOKEN
# Ex: just ds_tool -t textgen -d fixie-ai/boolq-audio -u fixie-ai/boolq-audio -c explanation
# Ex: just ds_tool -t textgen -d ylacombe/expresso -u fixie-ai/expresso -c continuation \
# -T "\"Continue the following sentence in a way that reflects a ‘{style}’ tone in a coherent style:\n{text}\""
# Ex: just ds_tool tts -d google/boolq -u fixie-ai/boolq-audio -c question -a audio --token $HF_WRITE_TOKEN
# Ex: just ds_tool textgen -d fixie-ai/boolq-audio -u fixie-ai/boolq-audio -c explanation
# Ex: just ds_tool textgen -d ylacombe/expresso -u fixie-ai/expresso -c continuation -T @expresso_template.txt
@dataclasses.dataclass
class DatasetToolArgs:
dataset_name: str = simple_parsing.field(alias="-d")
dataset_subset: str = simple_parsing.field(default="default", alias="-S")
dataset_subset: Optional[str] = simple_parsing.field(default=None, alias="-S")
dataset_split: Optional[str] = simple_parsing.field(default=None, alias="-s")

num_samples: Optional[int] = simple_parsing.field(default=None, alias="-n")
Expand All @@ -59,80 +90,38 @@ class DatasetToolArgs:
upload_name: Optional[str] = simple_parsing.field(default=None, alias="-u")
upload_branch: Optional[str] = simple_parsing.field(default="main", alias="-b")
num_shards: Optional[int] = simple_parsing.field(default=None, alias="-N")
private: bool = simple_parsing.field(default=False)

token: Optional[str] = None

task: Union[TtsTask, TextGenerationTask] = simple_parsing.subgroups(
{"tts": TtsTask, "textgen": TextGenerationTask}, # type: ignore
default_factory=TtsTask,
alias="-t",
positional=True,
)


def _tts_split(ds_split: datasets.Dataset, task: TtsTask, num_proc: int):
def _tts_sample(sample):
text = sample[task.column_name]
text = text["text"] if isinstance(text, dict) else text
sample[task.audio_column_name] = tts_client.tts(text)
return sample

print(f'TTS mapping "{task.column_name}" to "{task.audio_column_name}"...')

return ds_split.map(_tts_sample, num_proc=num_proc).cast_column(
task.audio_column_name, datasets.Audio(sampling_rate=task.sample_rate)
)


def _text_gen_split(
ds_split: datasets.Dataset, task: TextGenerationTask, num_proc: int
):
def _text_gen_sample(sample):
input_text = task.template.format(**sample)
response = chat_client.chat.completions.create(
model=task.language_model,
messages=[{"role": "user", "content": input_text}],
max_tokens=task.max_tokens,
temperature=task.temperature,
)
sample[task.new_column_name] = response.choices[0].message.content
return sample

print(
f'Text gen for column: "{task.new_column_name}" with template:\n{task.template}'
)
return ds_split.map(_text_gen_sample, num_proc=num_proc)


def main(args: DatasetToolArgs):
ds_name = args.dataset_name
print(f'Loading dataset "{ds_name}" for task {args.task}')
data_dict: datasets.DatasetDict = datasets.load_dataset(
ds_name, args.dataset_subset, split=args.dataset_split
)

if args.dataset_split:
data_dict = datasets.DatasetDict(**{args.dataset_split: data_dict})

for split, ds_split in data_dict.items():
print(f'Processing split "{split}"...')
if args.num_samples:
ds_split = ds_split.select(range(args.num_samples))

if isinstance(args.task, TtsTask):
data_dict[split] = _tts_split(
ds_split, args.task, num_proc=args.num_workers
)

elif isinstance(args.task, TextGenerationTask):
data_dict[split] = _text_gen_split(
ds_split, args.task, num_proc=args.num_workers
)
data_dict[split] = args.task.map_split(ds_split, args.num_workers)

token = args.token or os.environ.get("HF_TOKEN")
hub_args: Dict[str, Any] = {
"config_name": args.dataset_subset,
"config_name": args.dataset_subset or "default",
"token": token,
"revision": args.upload_branch,
"private": args.private,
}
if args.num_shards is not None:
hub_args["num_shards"] = {split: args.num_shards for split in data_dict.keys()}
Expand Down
Loading