-
Notifications
You must be signed in to change notification settings - Fork 525
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.
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
Support Remote and HF promptfiles in hf_generate script #786
Open
samhavens
wants to merge
13
commits into
main
Choose a base branch
from
remote-hf-generate
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
13 commits
Select commit
Hold shift + click to select a range
d05911e
support remote and hf promptfiles
samhavens 7f9be51
add back newline
samhavens 05c601d
grammar
samhavens 459bed0
fix remote file loading
samhavens be8cee2
Merge branch 'main' into remote-hf-generate
samhavens fff2596
fixes per daniel
samhavens eb8bb94
overwrite
samhavens c13d9fd
add truncation
samhavens 68f078e
lint
samhavens ada5f21
Merge branch 'main' into remote-hf-generate
samhavens 854b938
Merge branch 'main' into remote-hf-generate
samhavens 93e81fd
update cli help
samhavens 532a5db
Merge branch 'main' into remote-hf-generate
dakinggg File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -4,12 +4,16 @@ | |
import os | ||
from typing import List, Optional | ||
|
||
from composer.utils import get_file, parse_uri | ||
from datasets import load_dataset | ||
|
||
PROMPTFILE_PREFIX = 'file::' | ||
PROMPTDATASET_PREFIX = 'hf://' | ||
|
||
|
||
def load_prompts(prompts: List[str], | ||
prompt_delimiter: Optional[str] = None) -> List[str]: | ||
"""Loads a set of prompts, both free text and from file. | ||
"""Loads a set of prompts, both free text and from file or HF dataset. | ||
|
||
Args: | ||
prompts (List[str]): List of free text prompts and prompt files | ||
|
@@ -21,9 +25,16 @@ def load_prompts(prompts: List[str], | |
""" | ||
prompt_strings = [] | ||
for prompt in prompts: | ||
backend, _, _ = parse_uri(prompt) | ||
if prompt.startswith(PROMPTFILE_PREFIX): | ||
prompts = load_prompts_from_file(prompt, prompt_delimiter) | ||
prompt_strings.extend(prompts) | ||
elif prompt.startswith(PROMPTDATASET_PREFIX): | ||
prompts = load_prompts_from_dataset(prompt, prompt_delimiter) | ||
prompt_strings.extend(prompts) | ||
elif backend not in ['', None]: | ||
prompts = load_prompts_from_remote(prompt, prompt_delimiter) | ||
prompt_strings.extend(prompts) | ||
else: | ||
prompt_strings.append(prompt) | ||
return prompt_strings | ||
|
@@ -45,6 +56,7 @@ def load_prompts_from_file(prompt_path: str, | |
raise ValueError(f'prompt_path_str must start with {PROMPTFILE_PREFIX}') | ||
|
||
_, prompt_file_path = prompt_path.split(PROMPTFILE_PREFIX, maxsplit=1) | ||
# local file | ||
prompt_file_path = os.path.expanduser(prompt_file_path) | ||
if not os.path.isfile(prompt_file_path): | ||
raise FileNotFoundError( | ||
|
@@ -56,3 +68,73 @@ def load_prompts_from_file(prompt_path: str, | |
if prompt_delimiter is None: | ||
return [prompt_string] | ||
return [i for i in prompt_string.split(prompt_delimiter) if i] | ||
|
||
|
||
def load_prompts_from_remote(prompt_path: str, | ||
prompt_delimiter: Optional[str] = None | ||
) -> List[str]: | ||
"""Load a set of prompts from object storage. | ||
|
||
Args: | ||
prompt_path (str): Path for text file | ||
prompt_delimiter (Optional str): Delimiter for text file | ||
If not provided, assumes the prompt file is a single prompt (non-delimited) | ||
|
||
Returns: | ||
List of prompt string(s) | ||
""" | ||
backend, _, _ = parse_uri(prompt_path) | ||
if backend in ['', None]: | ||
raise ValueError( | ||
f'prompt_path_str must start with s3:// etc if using object storage' | ||
) | ||
|
||
local_path = prompt_path.split('/')[-1] | ||
get_file(path=prompt_path, destination=local_path, overwrite=True) | ||
|
||
with open(local_path, 'r') as f: | ||
prompt_string = f.read() | ||
|
||
if prompt_delimiter is None: | ||
return [prompt_string] | ||
return [i for i in prompt_string.split(prompt_delimiter) if i] | ||
|
||
|
||
def load_prompts_from_dataset(dataset_path: str, | ||
prompt_delimiter: Optional[str] = None | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'd prefer you just add a new argument to the script for the dataset column name instead of overloading the delimiter param |
||
) -> List[str]: | ||
"""Load a set of prompts from a huggingface dataset. | ||
|
||
Args: | ||
dataset_path (str): Path for dataset | ||
prompt_delimiter (Optional str): We misuse the delimiter here to specify | ||
the name of the prompt column in the dataset. If not provided, assumes the | ||
prompt column is named 'prompt'. | ||
|
||
Returns: | ||
List of prompt string(s) | ||
""" | ||
if not dataset_path.startswith(PROMPTDATASET_PREFIX): | ||
raise ValueError(f'dataset_path must start with {PROMPTDATASET_PREFIX}') | ||
|
||
_, dataset_path = dataset_path.split(PROMPTDATASET_PREFIX, maxsplit=1) | ||
|
||
try: | ||
dataset = load_dataset(dataset_path, token=True) | ||
except: | ||
dataset = load_dataset(dataset_path) | ||
|
||
prompt_strings = [] | ||
if prompt_delimiter is None: | ||
prompt_delimiter = 'prompt' | ||
try: | ||
ds = dataset['train'] | ||
except: | ||
ds = dataset | ||
|
||
if prompt_delimiter not in ds.column_names: | ||
raise ValueError(f'{prompt_delimiter} not in dataset columns.') | ||
for prompt in ds[prompt_delimiter]: | ||
prompt_strings.append(prompt) | ||
|
||
return prompt_strings |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
you probably want to make
local_path
a tmp location?