| import torch |
| from torch.utils.data import Dataset |
| from tokenizers import Tokenizer |
| from typing import List, Dict |
|
|
| from src.common.prompts import format_prompt, resolve_eos_id, resolve_pad_id |
|
|
| class SFTDataset(Dataset): |
| """ |
| Dataset for Supervised Fine-Tuning. |
| Formats Alpaca/Instruction data by masking the prompt loss and only computing loss on the response. |
| """ |
| def __init__(self, data: List[Dict[str, str]], tokenizer: Tokenizer, max_len: int = 512): |
| self.samples = [] |
|
|
| pad_token_id = resolve_pad_id(tokenizer) |
| eos_token_id = resolve_eos_id(tokenizer) |
|
|
| for item in data: |
| |
| prompt = format_prompt(item["prompt"]) |
| response = item["response"] |
|
|
| enc_prompt = tokenizer.encode(prompt).ids |
| enc_response = tokenizer.encode(response).ids |
| if eos_token_id is not None: |
| enc_response = enc_response + [eos_token_id] |
| |
| input_ids = enc_prompt + enc_response |
| |
| labels = [-100] * len(enc_prompt) + enc_response |
| |
| |
| if len(input_ids) > max_len: |
| input_ids = input_ids[:max_len] |
| labels = labels[:max_len] |
| else: |
| pad_len = max_len - len(input_ids) |
| input_ids = input_ids + [pad_token_id] * pad_len |
| labels = labels + [-100] * pad_len |
| |
| self.samples.append(( |
| torch.tensor(input_ids, dtype=torch.long), |
| torch.tensor(labels, dtype=torch.long) |
| )) |
|
|
| def __len__(self): |
| return len(self.samples) |
|
|
| def __getitem__(self, idx): |
| return self.samples[idx] |
|
|
