from transformers import Trainer, TrainingArguments, GPT2LMHeadModel, GPT2Tokenizer
import csv
import pandas as pd
import torch
import numpy as np
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler


class AnimalDataset(Dataset):
    def __init__(self, file_name):
        self.data = self._load_data(file_name)
        self.tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
        self.max_length = self._get_max_length()

    def _load_data(self, file_name):
        list_data = []
        with open(file_name, encoding='utf-8') as f:
            reader = csv.reader(f)
            headers = next(reader)
            for row in reader:
                for i in range(1, 4):
                    if row[i]:
                        q = f"What are the characteristics of {row[0]}?"
                        list_data.append([q, row[i]])
        return pd.DataFrame(data=list_data, columns=['Question', 'Description'])

    def _get_max_length(self):
        max_length = 0
        for index, row in self.data.iterrows():
            inputs = self.tokenizer.encode(row['Question'], add_special_tokens=True)
            labels = self.tokenizer.encode(row['Description'], add_special_tokens=True)
            max_length = max(max_length, len(inputs), len(labels))
        return max_length
    def _preprocess_sequence(self, sequence):
        if len(sequence) < self.max_length:
            sequence.extend([self.tokenizer.pad_token_id] * (self.max_length - len(sequence)))
        elif len(sequence) > self.max_length:
            sequence = sequence[:self.max_length]
        return sequence

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        question = self.data.iloc[index, self.data.columns.get_loc('Question')]
        description = self.data.iloc[index, self.data.columns.get_loc('Description')]

        inputs = self.tokenizer.encode(question, add_special_tokens=True)
        labels = self.tokenizer.encode(description, add_special_tokens=True)

        inputs = self._preprocess_sequence(inputs)
        labels = self._preprocess_sequence(labels)

        return {
            'input_ids': torch.tensor(inputs, dtype=torch.long),
            'labels': torch.tensor(labels, dtype=torch.long)
        }


def finetune(dataset_):
    model = GPT2LMHeadModel.from_pretrained("gpt2")
    training_args = TrainingArguments(
        output_dir='/content/sample_data/model1',
        num_train_epochs=1000,
        per_device_train_batch_size=1,
        per_device_eval_batch_size=1,
        logging_dir='/content/sample_data/log',
        learning_rate=1e-5,
        save_steps=False,
        remove_unused_columns=False
    )

    ds = AnimalDataset(dataset_)
     # 计算每个类别的样本权重
    labels = [item['labels'] for item in ds]
    class_weights = torch.tensor(1.0 / np.bincount(labels), dtype=torch.double)
    sample_weights = [class_weights[label].item() for label in labels]

    # 使用加权随机采样器来创建数据加载器
    sampler = WeightedRandomSampler(sample_weights, len(ds), replacement=True)
    dataloader = DataLoader(ds, batch_size=1, sampler=sampler)

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataloader,
        data_collator=lambda data: {'input_ids': torch.stack([item['input_ids'] for item in data]),
                                    'labels': torch.stack([item['labels'] for item in data])}
    )
    trainer.train()


df = "/content/sample_data/animal_description.csv"
finetune(df)

报错---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-27-7e820fe902b2> in <cell line: 95>()
     93 
     94 df = "/content/sample_data/animal_description.csv"
---> 95 finetune(df)

2 frames
<ipython-input-27-7e820fe902b2> in __getitem__(self, index)
     54 
     55         return {
---> 56             'input_ids': torch.tensor(inputs, dtype=torch.long),
     57             'labels': torch.tensor(labels, dtype=torch.long)
     58         }

TypeError: 'NoneType' object cannot be interpreted as an integer