|
import torch |
|
import torch.nn as nn |
|
from torch.utils.data import DataLoader, Dataset |
|
import tensorflow as tf |
|
import numpy as np |
|
|
|
def pad_sequences(sequences, max_seq_len=0): |
|
"""Pad sequences to max length in sequence.""" |
|
max_seq_len = max(max_seq_len, max(len(sequence) for sequence in sequences)) |
|
padded_sequences = np.zeros((len(sequences), max_seq_len)) |
|
for i, sequence in enumerate(sequences): |
|
padded_sequences[i][:len(sequence)] = sequence |
|
return padded_sequences |
|
|
|
class CustomDataSetManager(Dataset): |
|
def __init__(self, text_seq, line_num, total_line): |
|
self.text_seq = text_seq |
|
self.line_num_one_hot = line_num |
|
self.total_line_one_hot = total_line |
|
|
|
def __len__(self): |
|
return len(self.text_seq) |
|
|
|
def __str__(self): |
|
return f"<Dataset(N={len(self)})>" |
|
|
|
def __getitem__(self, index): |
|
X = self.text_seq[index] |
|
line_num = self.line_num_one_hot[index] |
|
total_line = self.total_line_one_hot[index] |
|
return [X, len(X), line_num, total_line] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def collate_fn(self, batch): |
|
"""Processing on a batch""" |
|
|
|
text_seq = [item[0] for item in batch] |
|
seq_lens = [item[1] for item in batch] |
|
line_nums = [item[2] for item in batch] |
|
total_lines = [item[3] for item in batch] |
|
|
|
|
|
pad_text_seq = pad_sequences(sequences=text_seq) |
|
|
|
|
|
line_nums = torch.tensor(tf.one_hot(line_nums, depth=20).numpy()) |
|
|
|
|
|
total_lines = torch.tensor(tf.one_hot(total_lines, depth=24).numpy()) |
|
|
|
|
|
pad_text_seq = torch.tensor(pad_text_seq, dtype=torch.long) |
|
seq_lens = torch.tensor(seq_lens, dtype=torch.long) |
|
|
|
return pad_text_seq, seq_lens, line_nums, total_lines |
|
|
|
def create_dataloader(self, batch_size, shuffle=False, drop_last=False): |
|
dataloader = DataLoader(dataset=self, batch_size=batch_size, collate_fn=self.collate_fn, shuffle=shuffle, drop_last=drop_last, pin_memory=True) |
|
return dataloader |
|
|