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