Spaces:

Ashmal
/

MobiLlama

Runtime error

App Files Files Community

MobiLlama / fastchat /train /train_flant5.py

Ashmal

Upload folder using huggingface_hub

5472531 verified 4 months ago

raw history blame contribute delete

No virus

15 kB

	# Adapted from tatsu-lab@stanford_alpaca. Below is the original copyright:
	# Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
	#
	# 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 collections import defaultdict
	import copy
	import os
	from dataclasses import dataclass, field
	import random
	import json
	import logging
	import pathlib
	from typing import Dict, Optional, Sequence

	import torch
	import torch.distributed as dist

	import transformers
	from torch.utils.data import Dataset
	from transformers import Trainer, AddedToken

	from fastchat.model.model_adapter import get_conversation_template

	default_conversation = get_conversation_template("t5")

	# TODO: import and use code from ../data/dataset.py

	IGNORE_INDEX = -100
	DEFAULT_PAD_TOKEN = "[PAD]"
	DEFAULT_EOS_TOKEN = "</s>"
	DEFAULT_BOS_TOKEN = "</s>"
	DEFAULT_UNK_TOKEN = "</s>"


	@dataclass
	class ModelArguments:
	model_name_or_path: Optional[str] = field(default="facebook/opt-125m")


	@dataclass
	class DataArguments:
	data_path: str = field(
	default=None, metadata={"help": "Path to the training data."}
	)
	lazy_preprocess: bool = False
	num_data: int = -1
	preprocessed_path: str = field(
	default=None, metadata={"help": "Path to the preprocessed training data."}
	)


	@dataclass
	class TrainingArguments(transformers.TrainingArguments):
	cache_dir: Optional[str] = field(default=None)
	optim: str = field(default="adamw_torch")
	model_max_length: int = field(
	default=2048,
	metadata={
	"help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
	},
	)


	def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
	"""Collects the state dict and dump to disk."""
	state_dict = trainer.model.state_dict()
	if trainer.args.should_save:
	cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
	del state_dict
	trainer._save(output_dir, state_dict=cpu_state_dict) # noqa


	def smart_tokenizer_and_embedding_resize(
	special_tokens_dict: Dict,
	other_tokens,
	tokenizer: transformers.PreTrainedTokenizer,
	model: transformers.PreTrainedModel,
	):
	"""Resize tokenizer and embedding.

	Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
	"""
	num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
	for new_token in other_tokens:
	num_new_tokens += tokenizer.add_tokens(AddedToken(new_token, normalized=False))

	model.resize_token_embeddings(len(tokenizer))

	if num_new_tokens > 0:
	input_embeddings = model.get_input_embeddings().weight.data
	output_embeddings = model.get_output_embeddings().weight.data

	input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(
	dim=0, keepdim=True
	)
	output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(
	dim=0, keepdim=True
	)

	input_embeddings[-num_new_tokens:] = input_embeddings_avg
	output_embeddings[-num_new_tokens:] = output_embeddings_avg


	def _tokenize_fn(
	strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer
	) -> Dict:
	"""Tokenize a list of strings."""
	tokenized_list = [
	tokenizer(
	text,
	return_tensors="pt",
	padding="longest",
	max_length=tokenizer.model_max_length,
	truncation=True,
	)
	for text in strings
	]
	input_ids = labels = [tokenized.input_ids[0] for tokenized in tokenized_list]
	input_ids_lens = labels_lens = [
	tokenized.input_ids.ne(tokenizer.pad_token_id).sum().item()
	for tokenized in tokenized_list
	]
	return dict(
	input_ids=input_ids,
	labels=labels,
	input_ids_lens=input_ids_lens,
	labels_lens=labels_lens,
	)


	def _form_qa(
	q_list,
	a_list,
	tokenized_conversation,
	tokenized_lens,
	speakers,
	header_len,
	max_length,
	eos_id,
	):
	cur_idx = header_len
	conv_len = len(tokenized_conversation)

	for tokenized_len, speaker in zip(tokenized_lens, speakers):
	if cur_idx >= conv_len:
	break
	if speaker == "gpt":
	# truncate answer if it is too long
	content_a = None
	if tokenized_len > max_length:
	content_a = tokenized_conversation[cur_idx : cur_idx + max_length]
	else:
	content_a = tokenized_conversation[cur_idx : cur_idx + tokenized_len]
	content_a.append(eos_id)
	a_list.append(content_a)
	content_q = None
	if cur_idx >= max_length:
	content_q = tokenized_conversation[cur_idx - max_length : cur_idx]
	else:
	content_q = tokenized_conversation[:cur_idx]
	content_q.append(eos_id)
	q_list.append(content_q)
	# asser the last token is actually a EOS for an answer
	assert a_list[-1][-1] == eos_id, "Last Token is not EOS!"
	cur_idx += tokenized_len


	def _add_speaker_and_signal(header, source, get_conversation=True):
	"""Add speaker and start/end signal on each round."""
	BEGIN_SIGNAL = "### "
	END_SIGNAL = "\n"
	conversation = header

	unknown_role = "unknown" # use default unknown role
	roles = {
	"human": default_conversation.roles[0], # human role
	"gpt": default_conversation.roles[1], # gpt role
	}

	for i in range(len(source)):
	sentence = source[i]
	sentence_from = sentence["from"].lower()

	# TODO(Dacheng): verify this is a good way to split sentences
	if sentence_from == "human":
	# if this is not the last sentence
	if i != len(source) - 1:
	next_sentence = source[i + 1]
	sentence["value"] = (
	BEGIN_SIGNAL
	+ roles.get(sentence_from, unknown_role)
	+ ": "
	+ sentence["value"]
	+ END_SIGNAL
	+ BEGIN_SIGNAL
	+ roles.get(next_sentence["from"].lower(), unknown_role)
	+ ": "
	)
	else:
	# if human is the last speaker, it does not contribute to an answer
	pass
	else:
	sentence["value"] = sentence["value"] + END_SIGNAL
	if get_conversation:
	conversation += sentence["value"]

	return conversation


	def preprocess(
	sources: Sequence[str],
	tokenizer: transformers.PreTrainedTokenizer,
	) -> Dict:
	"""
	Given a list of sources, each is a conversation list. This transform:
	1. Add signal '### ' at the beginning each sentence, with end signal '\n';
	2. Concatenate conversations together;
	3. Tokenize the concatenated conversation;
	4. Make a deepcopy as the target. Mask human words with IGNORE_INDEX.
	"""
	# add end signal and concatenate together
	conversations = []
	header = f"{default_conversation.system_message}\n\n"
	for source in sources:
	conversation = _add_speaker_and_signal(header, source, tokenizer)
	conversations.append(conversation)
	# TODO(Dacheng): This is related to whether the dataset has been truncated..
	# Assume we get long conversations, don't pad, don't return tensor
	tokenized_conversations = tokenizer(conversations, max_length=None)["input_ids"]
	q_list = []
	a_list = []
	# count for EOS length
	header_len = _tokenize_fn([header], tokenizer)["input_ids_lens"][0] - 1
	from tqdm import tqdm

	for tokenized_conversation, source in tqdm(zip(tokenized_conversations, sources)):
	tokenized_sentence = _tokenize_fn([s["value"] for s in source], tokenizer)
	tokenized_lens = tokenized_sentence["input_ids_lens"]
	tokenized_lens = [l - 1 for l in tokenized_lens]
	speakers = [sentence["from"] for sentence in source]
	ids = tokenized_sentence["input_ids"]
	_form_qa(
	q_list,
	a_list,
	tokenized_conversation,
	tokenized_lens,
	speakers,
	header_len,
	tokenizer.model_max_length,
	tokenizer.eos_token_id,
	)
	return dict(input_ids=q_list, labels=a_list)


	class SupervisedDataset(Dataset):
	"""Dataset for supervised fine-tuning."""

	def __init__(
	self,
	data_path: str,
	tokenizer: transformers.PreTrainedTokenizer,
	preprocessed_path,
	num_data,
	):
	super(SupervisedDataset, self).__init__()

	# save to file
	# Make sure only the first process is processing the dataset
	if dist.get_rank() != 0:
	dist.barrier()
	self.preprocessed_path = preprocessed_path
	if os.path.exists(self.preprocessed_path):
	logging.warning("loading from preprocessed data")
	with open(self.preprocessed_path, "r") as f:
	data_dict = json.load(f)
	if dist.get_rank() == 0:
	dist.barrier()
	else:
	if not os.path.exists("preprocessed_data"):
	os.mkdir("preprocessed_data")
	assert dist.get_rank() == 0, "Only the first process should process"
	logging.warning("Loading data...")
	list_data_dict = json.load(open(data_path, "r"))

	logging.warning("Formatting inputs...")
	sources = []

	sources = [example["conversations"] for example in list_data_dict]

	data_dict = preprocess(sources, tokenizer)
	json_data_dict = json.dumps(data_dict)

	# Remember to close file to avoid concurrent r/w
	with open(self.preprocessed_path, "w") as f:
	f.write(json_data_dict)

	# Release barrier
	dist.barrier()

	if num_data != -1:
	data_dict["input_ids"] = data_dict["input_ids"][:num_data]
	data_dict["labels"] = data_dict["labels"][:num_data]

	# Shuffle data to see more conversations, if only train on partial data
	temp = list(zip(data_dict["input_ids"], data_dict["labels"]))
	random.shuffle(temp)
	res1, res2 = zip(*temp)
	data_dict["input_ids"], data_dict["labels"] = list(res1), list(res2)

	# Dacheng: Get rid of short QA pair
	self.input_ids = copy.deepcopy(data_dict["input_ids"])
	self.labels = copy.deepcopy(data_dict["labels"])
	length_arr = defaultdict(int)
	for idx, (input, label) in enumerate(
	zip(data_dict["input_ids"], data_dict["labels"])
	):
	length_arr[str(len(label) // 100)] += 1
	if len(input) <= 5:
	del_idx = self.input_ids.index(input)
	self.input_ids.pop(del_idx)
	self.labels.pop(del_idx)
	if len(label) <= 5:
	del_idx = self.labels.index(label)
	self.input_ids.pop(del_idx)
	self.labels.pop(del_idx)

	for input, label in zip(self.input_ids, self.labels):
	assert len(input) >= 5
	assert len(label) >= 5

	def __len__(self):
	return len(self.input_ids)

	def __getitem__(self, i) -> Dict[str, torch.Tensor]:
	return dict(input_ids=self.input_ids[i], labels=self.labels[i])


	@dataclass
	class DataCollatorForSupervisedDataset(object):
	"""Collate examples for supervised fine-tuning."""

	tokenizer: transformers.PreTrainedTokenizer

	def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
	input_ids, labels = tuple(
	[
	torch.as_tensor(instance[key], dtype=torch.int64)
	for instance in instances
	]
	for key in ("input_ids", "labels")
	)
	input_ids = torch.nn.utils.rnn.pad_sequence(
	input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
	)
	labels = torch.nn.utils.rnn.pad_sequence(
	labels, batch_first=True, padding_value=IGNORE_INDEX
	)
	ret = dict(
	input_ids=input_ids,
	labels=labels,
	attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
	)
	torch.set_printoptions(profile="full")
	return ret


	def make_supervised_data_module(
	tokenizer: transformers.PreTrainedTokenizer, data_args
	) -> Dict:
	"""Make dataset and collator for supervised fine-tuning."""
	dataset_cls = SupervisedDataset
	train_dataset = dataset_cls(
	tokenizer=tokenizer,
	data_path=data_args.data_path,
	preprocessed_path=data_args.preprocessed_path,
	num_data=data_args.num_data,
	)
	data_collator = DataCollatorForSupervisedDataset(tokenizer=tokenizer)
	return dict(
	train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator
	)


	def train():
	parser = transformers.HfArgumentParser(
	(ModelArguments, DataArguments, TrainingArguments)
	)
	model_args, data_args, training_args = parser.parse_args_into_dataclasses()

	model = transformers.AutoModelForSeq2SeqLM.from_pretrained(
	model_args.model_name_or_path,
	cache_dir=training_args.cache_dir,
	)
	# Dacheng: Note we can only use T5Tokenizer, otherwise it will prepend
	# a space before special tokens.
	tokenizer = transformers.T5Tokenizer.from_pretrained(
	model_args.model_name_or_path,
	cache_dir=training_args.cache_dir,
	model_max_length=training_args.model_max_length,
	padding_side="right",
	use_fast=False,
	)

	smart_tokenizer_and_embedding_resize(
	special_tokens_dict=dict(pad_token=DEFAULT_PAD_TOKEN),
	other_tokens=["<", "{", "\n", "}", "`", " ", "\\", "^", "\t"],
	tokenizer=tokenizer,
	model=model,
	)

	data_module = make_supervised_data_module(tokenizer=tokenizer, data_args=data_args)
	trainer = Trainer(
	model=model, tokenizer=tokenizer, args=training_args, **data_module
	)

	if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
	trainer.train(resume_from_checkpoint=True)
	else:
	trainer.train()
	trainer.save_state()
	safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)


	if __name__ == "__main__":
	train()