github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-20 09:01:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: fix sft (#3568)

Browse Source
This commit is contained in:
tingfeng cao
2023-04-17 16:47:44 +08:00
committed by GitHub
parent 6e7e43c6fe
commit 7788e0b0a5
3 changed files with 13 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
applications/Chat/coati/dataset/sft_dataset.py
View File

@@ -53,29 +53,25 @@ class SFTDataset(Dataset):
def __init__(self, dataset, tokenizer: Callable, max_length: int = 512) -> None:
super().__init__()
# self.prompts = []
self.input_ids = []
for data in tqdm(dataset, disable=not is_rank_0()):
prompt = data['prompt'] + data['completion'] + "<|endoftext|>"
prompt = data['prompt'] + data['completion'] + tokenizer.eos_token
prompt_token = tokenizer(prompt,
max_length=max_length,
padding="max_length",
truncation=True,
return_tensors="pt")
# self.prompts.append(prompt_token)s
self.input_ids.append(prompt_token)
self.labels = copy.deepcopy(self.input_ids)
self.input_ids.append(prompt_token['input_ids'][0])
self.labels = copy.deepcopy(self.input_ids)
def __len__(self):
length = len(self.prompts)
length = len(self.input_ids)
return length
def __getitem__(self, idx):
# dict(input_ids=self.input_ids[i], labels=self.labels[i])
return dict(input_ids=self.input_ids[idx], labels=self.labels[idx])
# return dict(self.prompts[idx], self.prompts[idx])
def _tokenize_fn(strings: Sequence[str], tokenizer: transformers.PreTrainedTokenizer, max_length: int) -> Dict: