import pandas as pd
import numpy as np
from torch.utils.data import Dataset, DataLoader

class WikiDataset(Dataset):
    def __init__(self, data_path, transform=None):
        self.data_frame = pd.read_csv(data_path)
        self.data_frame = self.data_frame.dropna()

        self.transform = transform
    
    def __len__(self):
        return len(self.data_frame)
    
    def __getitem__(self, index):

        data_at_index = self.data_frame.iloc[index]
        content = data_at_index['content']
        label = data_at_index['label']
        
        if self.transform:
            content = self.transform(content)
        
        return content, label

    def get_data_loader(self, batch_size=32, shuffle=False):
        return DataLoader(self, batch_size=batch_size, shuffle=shuffle)