File size: 2,967 Bytes
2c15172 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | import json
import torch
from typing import Final, Dict
from torch.utils.data import Dataset
class FriendsDataset(Dataset):
""" Dataset class for Friends transcript dialogue data (.json).
Args:
data: the .json corpus file (path name)
tokenizer: custom GPT-2 tokenizer
maxt: maximum token length, default 128
"""
# hard code speakers
SPEAKER_LOOKUP: Final[Dict[str, int]] = {"ROSS": 0, "MONICA": 1, "CHANDLER": 2,
"JOEY": 3, "RACHEL": 4, "PHOEBE": 5,
"GUNTHER": 6, "JANICE": 7, "RICHARD": 8,
"CAROL": 9, "SUSAN": 10, "MIKE": 11, "OTHER": 12}
def __init__(self, data, tokenizer, maxt = 128):
self.sequences = []
self.responders = []
self.tokenizer = tokenizer
self.maxt = maxt
self.pad_id = tokenizer.pad_token_id
# import corpus
with open(data, 'r', encoding = 'utf-8') as f:
corpus = json.load(f)
def format_turn(turn: dict[str, str]) -> str:
""" Context Format Helper
"""
speaker = turn['speaker'].upper()
return f"<SPEAKER={speaker}> {turn['text']}", speaker
# unravel scenes and construct input sequence tensors
for scene in corpus:
turns = scene['turns']
n = len(turns)
# Context Windows (1,2,3)
for i in range(n):
for window in range(1, 4):
if i - window < 0: continue
context = turns[i - window:i]
response = turns[i]
# build sequence
context_str = "\n".join(format_turn(t)[0] for t in context)
response_str, responder = format_turn(response)
sequence = ("<CONTEXT>\n" + context_str + \
"\n</CONTEXT>\n<RESPONSE>\n" + \
response_str + "\n<EOT>")
self.sequences.append(sequence)
self.responders.append(self.SPEAKER_LOOKUP.get(
responder, self.SPEAKER_LOOKUP["OTHER"]))
def __len__(self): return len(self.sequences)
def __getitem__(self, index):
""" Tokenization on the fly
"""
encoded = self.tokenizer(self.sequences[index], add_special_tokens = False,
max_length = self.maxt,
truncation = True,
padding = 'max_length',
return_tensors = 'pt')
return {'input_ids': encoded['input_ids'].squeeze(0),
'attention_mask': encoded['attention_mask'].squeeze(0),
'responder': torch.tensor(self.responders[index], dtype = torch.long)}
|