import torch 

from transformers import AutoTokenizer,AutoModel, AdamW,BertForSequenceClassification

from torch.utils.data import DataLoader, Dataset, random_split 

import pandas as pd 

from tqdm import tqdm 

import random 

from transformers import BertForSequenceClassification

# 定义一个自定义数据集类，用于加载情感分析数据集
class SentimentDataset(Dataset):
    def __init__(self, dataframe, tokenizer, max_length=128):
        self.dataframe = dataframe  # 存储数据帧
        self.tokenizer = tokenizer  # 存储tokenizer对象
        self.max_length = max_length  # 存储最大长度参数

    # 返回数据集中的样本数量
    def __len__(self):
        return len(self.dataframe)

    # 根据给定索引返回对应的数据样本
    def __getitem__(self, idx):
        text = self.dataframe.iloc[idx]['review']  # 获取指定索引处的文本数据
        label = 1 if self.dataframe.iloc[idx]['label'] == 1 else 0  # 获取指定索引处的标签，并将其转换为0或1
        # 使用tokenizer对文本进行处理，并返回模型所需的输入格式
        encoding = self.tokenizer(text, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='pt')
        return {
            'input_ids': encoding['input_ids'].flatten(),  # 输入token的ID
            'attention_mask': encoding['attention_mask'].flatten(),  # 注意力掩码
            'labels': torch.tensor(label, dtype=torch.long)  # 标签
        }

model.train()

for epoch in range(3):
    for batch in tqdm(train_loader, desc="Epoch {}".format(epoch + 1)):
        # 获取输入的ID
        input_ids = batch['input_ids'].to(device)
        # 获取注意力掩码
        attention_mask = batch['attention_mask'].to(device)
        # 获取标签
        labels = batch['labels'].to(device)

        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

import torch
from torch.utils.data import Dataset

# 定义用于加载测试数据集的类
class TestDataset(Dataset):
    def __init__(self, file_path, tokenizer, max_length=128):
        self.data = []
        with open(file_path, "r", encoding="utf-8") as file:
            for line in file:
                label, text = line.strip().split("\t") 
                self.data.append((int(label), text))  # 添加到数据列表中
        self.tokenizer = tokenizer  # 存储分词器对象
        self.max_length = max_length  # 存储最大长度参数

    def __len__(self):
        return len(self.data)  # 返回数据集的样本数量

    def __getitem__(self, idx):
        label, text = self.data[idx]  # 获取指定索引处的标签和文本数据
        # 使用tokenizer处理文本，并返回模型所需的输入格式
        encoding = self.tokenizer(text, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='pt')
        return {
            'input_ids': encoding['input_ids'].flatten(),  # 输入token的ID
            'attention_mask': encoding['attention_mask'].flatten(),  # 注意力掩码
            'labels': torch.tensor(label, dtype=torch.long)  # 标签
        }

# 加载测试数据集
test_dataset = TestDataset("weibo-emo/test.txt", tokenizer, max_length=128)  #=
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=8, shuffle=False) 

# 评估模型
model.eval() 
total_eval_accuracy = 0
for batch in tqdm(test_loader, desc="Evaluating"):  
    input_ids = batch['input_ids'].to(device) 
    attention_mask = batch['attention_mask'].to(device)  
    labels = batch['labels'].to(device)  

    with torch.no_grad():  
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels) 
    
    logits = outputs.logits  # 获取预测的logits
    preds = torch.argmax(logits, dim=1)  # 获取预测结果
    accuracy = (preds == labels).float().mean() 
    total_eval_accuracy += accuracy.item()  

average_eval_accuracy = total_eval_accuracy / len(test_loader)  # 计算平均准确率
print("Test Accuracy:", average_eval_accuracy)