File size: 631 Bytes
fbf7e95 |
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 |
from torch.utils.data import Dataset
import numpy as np
import pandas as pd
class SimilarityVectorDataset(Dataset):
def __init__(self, processed_path: str, transform=None):
self.transform = transform
self.data = pd.read_csv(processed_path)
def __len__(self):
return self.data.shape[0]
def __getitem__(self, idx):
row = self.data.iloc[idx].to_dict()
label = float(float(row['cid']) == 1.0)
if self.transform:
row = self.transform(row)
row = np.array(list(row.values())).astype(float)
return row, label
|