from torch.utils.data import Dataset
from torchtext.data.utils import get_tokenizer
import pandas as pd

#第一步的数据集的构造
#无论是计算机视觉任务还是NLP任务,都要继承Dataset
#构造一个用于读取数据的数据集
class NewsDataset(Dataset):
    # data对象中保存了原始数据,包括标签label和新闻描述text
    # data[0] = (3,"Wall St. Bears Claw Back Into the ...")
    def __init__(self,is_train=True) -> None:
        super().__init__()
        if is_train:
            data = pd.read_csv("./text_cls/ag_news_csv/train.csv")
        else:
            data = pd.read_csv("./text_cls/ag_news_csv/test.csv")
        self.examples = list()
        #返回一个内置的英文分词器
        tokenizer = get_tokenizer("basic_english")
        for index, row in data.iterrows():
            label =row.iloc[0]
            text = row.iloc[2]
            #将文本转为小写,并使用tokenizer进行分词
            tokenized_text = [token.lower() for token in tokenizer(text)]
            #将处理好的文本和标签,保存到list中
            self.examples.append((tokenized_text,label))
    
    def __len__(self):
        return len(self.examples)
    def __getitem__(self, index):
        return self.examples[index]

if __name__== "__main__":
    dataset = NewsDataset()
    print(len(dataset))
    print(dataset[0])