|
"""Implementation derived from https://github.com/tloen/alpaca-lora""" |
|
import os |
|
import sys |
|
from pathlib import Path |
|
|
|
|
|
wd = Path(__file__).parent.parent.resolve() |
|
sys.path.append(str(wd)) |
|
|
|
import torch |
|
import requests |
|
import json |
|
from torch.utils.data import random_split |
|
sys.path.append(os.getcwd()) |
|
from lit_llama.tokenizer import Tokenizer |
|
from tqdm import tqdm |
|
import numpy as np |
|
|
|
from options import option |
|
|
|
IGNORE_INDEX = -1 |
|
|
|
def prepare( |
|
destination_path: Path = Path("./data"), |
|
tokenizer_path: Path = Path("./checkpoints/lit-llama/tokenizer.model"), |
|
max_seq_length: int = 2560, |
|
seed: int = 42, |
|
mask_inputs: bool = False, |
|
split: str = "train" |
|
): |
|
"""Prepare the Alpaca dataset for instruction tuning. |
|
The output is a training and validation dataset saved as `train.pt` and `val.pt`, |
|
which stores the preprocessed and tokenized prompts and labels. |
|
""" |
|
|
|
destination_path.mkdir(parents=True, exist_ok=True) |
|
|
|
file_path = f'/comp_robot/lushunlin/MotionGPT/data/video_dataset/{split}.json' |
|
|
|
|
|
tokenizer = Tokenizer(tokenizer_path) |
|
|
|
with open(file_path, "r") as file: |
|
data = json.load(file) |
|
data_set = list(data) |
|
|
|
print(f"{split} set has {len(data_set):,} samples") |
|
|
|
print(f"Processing {split} split ...") |
|
data_set_new = [] |
|
for sample in tqdm(data_set): |
|
|
|
data_set_new.append(prepare_sample(sample, tokenizer, max_seq_length, mask_inputs)) |
|
|
|
|
|
data_set = data_set_new |
|
|
|
save_pt = f'/comp_robot/lushunlin/MotionGPT/data/video_dataset/{split}.pt' |
|
torch.save(data_set, save_pt) |
|
|
|
|
|
def prepare_sample(example: dict, tokenizer: Tokenizer, max_length: int, mask_inputs: bool = True): |
|
"""Processes a single sample. |
|
Each sample in the dataset consists of: |
|
- instruction: A string describing the task |
|
- input: A string holding a special input value for the instruction. |
|
This only applies to some samples, and in others this is empty. |
|
- output: The response string |
|
|
|
This function processes this data to produce a prompt text and a label for |
|
supervised training. The prompt text is formed as a single message including both |
|
the instruction and the input. The label/target is the same message but with the |
|
response attached. |
|
|
|
Finally, both the prompt and the label get tokenized. If desired, all tokens |
|
in the label that correspond to the original input prompt get masked out (default). |
|
|
|
|
|
""" |
|
|
|
|
|
|
|
full_prompt = generate_prompt_mlp(example) |
|
full_prompt_and_response = full_prompt + example['output'] |
|
|
|
encoded_full_prompt = tokenize(tokenizer, full_prompt, max_length=max_length, eos=False) |
|
encoded_full_prompt_and_response = tokenize(tokenizer, full_prompt_and_response, eos=True, max_length=max_length) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
labels = encoded_full_prompt_and_response.clone() |
|
if mask_inputs: |
|
labels[:len(encoded_full_prompt)] = IGNORE_INDEX |
|
|
|
|
|
|
|
return {**example, "sys_command": generate_system_command(), "input_ids": encoded_full_prompt_and_response, "input_ids_no_response": encoded_full_prompt, "labels": labels} |
|
|
|
|
|
def tokenize(tokenizer: Tokenizer, string: str, max_length: int, eos=True) -> torch.Tensor: |
|
return tokenizer.encode(string, bos=True, eos=eos, max_length=max_length) |
|
|
|
def detokenizer(tokenizer: Tokenizer, tensor: torch.Tensor): |
|
''' |
|
tokenizer.decode(torch.tensor([13866, 338])) |
|
''' |
|
return tokenizer.decode(tensor) |
|
|
|
|
|
def generate_prompt_mlp(example): |
|
"""Generates a standardized message to prompt the model with an instruction, optional input and a |
|
'response' field.""" |
|
|
|
|
|
|
|
|
|
|
|
if example["input"]: |
|
return ( |
|
f"A chat between a curious user and an artificial intelligence assistant, paired with an input that provides further context. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {example['instruction']} INPUT_VIDEO: {example['input']}. \nASSISTANT: " |
|
) |
|
return ( |
|
f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: {example['instruction']} ASSISTANT: " |
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_system_command(): |
|
return ( |
|
f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. " |
|
) |
|
|
|
|
|
def main(): |
|
args = option.get_args_parser() |
|
|
|
|
|
prepare(split='train_intern_human_2M_stage1_caption') |
|
prepare(split='val_intern_human_2M_stage1_caption') |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|