| import json |
| import torch |
| from torch.utils.data import Dataset |
| from datasets import load_dataset, Features, Value |
|
|
| from dataset.common import pre_processing_chat, post_processing_chat |
|
|
|
|
| class SFTDataset(Dataset): |
| def __init__(self, jsonl_path, tokenizer, max_length=1024): |
| super().__init__() |
| self.tokenizer = tokenizer |
| self.max_length = max_length |
| features = Features({'conversations': [{'role': Value('string'), 'content': Value('string'), 'reasoning_content': Value('string'), 'tools': Value('string'), 'tool_calls': Value('string')}]}) |
| self.samples = load_dataset('json', data_files=jsonl_path, split='train', features=features) |
| self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant\n', add_special_tokens=False).input_ids |
| self.eos_id = tokenizer(f'{tokenizer.eos_token}\n', add_special_tokens=False).input_ids |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def create_chat_prompt(self, conversations): |
| messages = [] |
| tools = None |
| for message in conversations: |
| message = dict(message) |
| if message.get("role") == "system" and message.get("tools"): |
| tools = json.loads(message["tools"]) if isinstance(message["tools"], str) else message["tools"] |
| if message.get("tool_calls") and isinstance(message["tool_calls"], str): |
| message["tool_calls"] = json.loads(message["tool_calls"]) |
| messages.append(message) |
| return self.tokenizer.apply_chat_template( |
| messages, |
| tokenize=False, |
| add_generation_prompt=False, |
| tools=tools |
| ) |
|
|
| def generate_labels(self, input_ids): |
| labels = [-100] * len(input_ids) |
| i = 0 |
| while i < len(input_ids): |
| if input_ids[i:i + len(self.bos_id)] == self.bos_id: |
| start = i + len(self.bos_id) |
| end = start |
| while end < len(input_ids): |
| if input_ids[end:end + len(self.eos_id)] == self.eos_id: |
| break |
| end += 1 |
| for j in range(start, min(end + len(self.eos_id), self.max_length)): |
| labels[j] = input_ids[j] |
| i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids) |
| else: |
| i += 1 |
| return labels |
|
|
| def __getitem__(self, index): |
| sample = self.samples[index] |
| conversations = pre_processing_chat(sample['conversations']) |
| prompt = self.create_chat_prompt(conversations) |
| prompt = post_processing_chat(prompt) |
| input_ids = self.tokenizer(prompt).input_ids[:self.max_length] |
| input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) |
| labels = self.generate_labels(input_ids) |
| return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long) |
|
|
