01bb07dba8
FocalLoRA, Should-It-Be-Executed-Or-Processed, and topicattack were nested git repos (with an inner FocalLoRA/data/FocalLoRA/.git as well). Drop their .git history and track the contents directly in this repo instead of as submodules/gitlinks. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
65 lines
2.6 KiB
Python
65 lines
2.6 KiB
Python
import os.path
|
|
|
|
import fire
|
|
import pandas as pd
|
|
|
|
|
|
def main(data_path: str, model_type: str, out_dir: str) -> None:
|
|
"""
|
|
Process the input data, generate text based on the specified model type, and save the
|
|
resulting datasets into training and testing sets.
|
|
|
|
Parameters:
|
|
data_path (str): The path to the input JSON file containing the data.
|
|
model_type (str): The type of model to be used for generating text. This should be one
|
|
of the keys in the `map_funcs` dictionary.
|
|
out_dir (str): The directory where the output JSON files will be saved.
|
|
|
|
The function performs the following steps:
|
|
1. Reads the input JSON file into a pandas DataFrame.
|
|
2. Retains only the 'inputs', 'goal_safe', and 'goal_unsafe' columns.
|
|
3. Extracts 'prompt_1' and 'prompt_2' from the 'inputs' column and removes 'inputs'.
|
|
4. Applies a text generation function based on the specified `model_type`.
|
|
5. Splits the data into training (80%) and testing (20%) sets.
|
|
6. Prints two random samples from the training set.
|
|
7. Saves the training and testing sets as JSON files in the specified output directory.
|
|
|
|
Raises:
|
|
KeyError: If `model_type` is not one of the predefined keys in the `map_funcs` dictionary.
|
|
"""
|
|
df = pd.read_json(data_path)
|
|
# remove columns except prompt_1 prompt_2 goal_safe goal_unsafe
|
|
df = df[['inputs', 'goal_safe', 'goal_unsafe']]
|
|
df['prompt_1'] = df['inputs'].apply(lambda x: x['prompt_1'])
|
|
df['prompt_2'] = df['inputs'].apply(lambda x: x['prompt_2'])
|
|
df = df.drop(columns=['inputs'])
|
|
|
|
map_funcs = {
|
|
"gemma-1.1": lambda x: x['prompt_1'] + x['goal_safe'] + "<end_of_turn>",
|
|
"Starling-LM": lambda x: x['prompt_1'] + x['goal_safe'] + "<|end_of_turn|>",
|
|
"Llama-3": lambda x: x['prompt_1'] + x['goal_safe'] + "<|eot_id|>",
|
|
"Llama-2": lambda x: x['prompt_1'] + x['goal_safe'] + "</s>",
|
|
"zephyr": lambda x: x['prompt_1'] + x['goal_safe'] + "</s>",
|
|
"Phi-3": lambda x: x['prompt_1'] + x['goal_safe'] + "<|end|>",
|
|
}
|
|
|
|
df['text'] = df.apply(map_funcs[model_type], axis=1)
|
|
|
|
df_train = df.sample(frac=0.8, random_state=42)
|
|
df_test = df.drop(df_train.index)
|
|
|
|
#print a few random samples
|
|
print(df_train['text'].sample(2).to_list())
|
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
df_train.to_json(f'{out_dir}/train_dataset.json',
|
|
orient='records',
|
|
force_ascii=False)
|
|
df_test.to_json(f'{out_dir}/test_dataset.json',
|
|
orient='records',
|
|
force_ascii=False)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
fire.Fire(main)
|