###main.py###
import torch.nn as nn
import torch
import random
import numpy  as np
import os

from model_utils.data import get_dataloader
from model_utils.model import myBertModel
from model_utils.train import train_val

# 固定随机种子函数
# 机器学习实验必须固定种子，否则每次运行的随机初始化（模型参数、数据打乱、优化器更新）会导致结果不一致，无法判断模型改进是否有效；
# seed=1 是常用的固定值，也可设为 42、1024 等，只要固定即可。
def seed_everything(seed):
    # 固定PyTorch CPU随机种子
    torch.manual_seed(seed)
    # 固定单个GPU随机种子
    torch.cuda.manual_seed(seed)
    # 固定所有GPU随机种子（多卡训练时用）
    torch.cuda.manual_seed_all(seed)
    # 关闭cudnn自动优化（避免不同运行的算法选择不同导致结果不一致）
    torch.backends.cudnn.benchmark = False
    # 强制cudnn使用确定性算法（保证结果可复现）
    torch.backends.cudnn.deterministic = True
    # 固定Python原生随机种子
    random.seed(seed)
    # 固定NumPy随机种子
    np.random.seed(seed)
    # 固定Python哈希种子（避免字典/集合等哈希操作的随机性）
    os.environ['PYTHONHASHSEED'] = str(seed)

# 超参数配置（核心可调整项）
#
#
model_name = 'MyModel'
num_class = 2       # 分类任务的类别数（二分类，对应酒店评论的正面/负面）
batchSize = 32      # 批量大小（需根据GPU显存调整，6GB显存建议设为8/16，32可能显存不足）
learning_rate = 0.0001      # 学习率（BERT微调的经典学习率，远小于普通CNN/RNN的1e-3）
loss = nn.CrossEntropyLoss()    # 损失函数（二分类/多分类的标准损失，自动处理one-hot标签）
epoch = 3       # 训练轮数（小数据集设3-5轮，避免过拟合）
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 数据文件路径（酒店评论数据）
data_path = "jiudian.txt"
# BERT预训练模型路径（使用huggingface的中文BERT-Base）
bert_path = 'bert-base-chinese'
# 模型保存路径（训练过程中保存最优模型）
save_path = 'model_save/'
seed_everything(1)

# 数据加载与模型初始化
##########################################
# 加载训练/验证DataLoader（batchsize=32）
train_loader, val_loader = get_dataloader(data_path, batchsize=batchSize)

# 初始化自定义BERT模型：传入预训练路径、类别数、设备，移到指定设备（GPU/CPU）
model = myBertModel(bert_path, num_class, device).to(device)    #由于报错说数据不在同一设备，所以·这里传入设备保证统一
# tokenizer, model = build_model_and_tokenizer(bert_path)


# 优化器与学习率调度器配置
# （无用代码）获取模型参数列表（仅定义未使用）
param_optimizer = list(model.parameters())

# 定义优化器：AdamW（BERT微调的标准优化器，带权重衰减）
# lr=学习率，weight_decay=权重衰减（L2正则，防止过拟合）
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate,weight_decay=0.0001)

# 学习率调度器：余弦退火+预热重启（CosineAnnealingWarmRestarts）
# T_0=20：每20个epoch重启一次学习率，eta_min=1e-9：学习率最低降到1e-9
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,T_0=20,eta_min=1e-9)



# 数据加载与模型初始化
#
#
trainpara = {'model': model,
             'train_loader': train_loader,
             'val_loader': val_loader,
             'scheduler': scheduler,
             'optimizer': optimizer,
             'learning_rate': learning_rate,
             'warmup_ratio' : 0.1,      #  预热比例（0.1）：BERT 微调经典配置，前 10% 的 step 学习率线性上升，避免初始震荡
             'weight_decay' : 0.0001,   #  权重衰减：和优化器中的配置一致，若train_val()内需要重新初始化优化器则会用到
             'use_lookahead' : True,    #是否启用 Lookahead 优化：对 AdamW 做改进，提升训练稳定性（可选功能）
             'loss': loss,
             'epoch': epoch,
             'device': device,
             'save_path': save_path,
             'max_acc': 0.85,           #验证集准确率阈值：初始设 0.85，只有超过该值才保存模型，避免保存效果差的模型
             'val_epoch' : 1    #训练多少轮验证一次
             }
train_val(trainpara)

###model.py###

import torch        #导入 PyTorch 核心库
import torch.nn as nn       #导入 PyTorch 神经网络模块，并简写为nn
import numpy as np          #处理数值计算（比如数据预处理、数组转换），虽然当前代码中没直接用，但属于机器学习标配（后续扩展会用到）
# from timm.models.vision_transformer import PatchEmbed, Block
# import torchvision.models as models     #torchvision是处理图像的库（比如 ResNet、VGG）,完全用不到，可直接删除`
import transformers         #作为下面具体类的 “父库”，确保后续导入的 BERT 相关类能正常调用
from transformers import (BertPreTrainedModel, BertConfig,
                          BertForSequenceClassification, BertTokenizer,BertModel,
                          )
# BertPreTrainedModel封装了 BERT 模型的通用功能（比如加载预训练权重、初始化参数、配置设备），BertModel、BertForSequenceClassification都是继承自它。
# BertConfig相当于 BERT 模型的 “配置文件”，告诉模型 “该用多大的隐藏层、多少个注意力头、词表有多大”。
# BertForSequenceClassification相当于把BertModel（特征提取）+ nn.Linear（分类头）打包好了，直接传入文本就能输出分类结果。相比BertModel不需要手动分类头
# BertTokenizer把中文文本（比如 “这家酒店很好”）转换成 BERT 能看懂的数字编码（input_ids）、注意力掩码（attention_mask）等。
# BertModel是BERT 的核心部分，接收分词后的输入，输出文本的特征向量（768 维），是整个模型的 “特征提取器”。



class myBertModel(nn.Module):
    def __init__(self, bert_path, num_class, device):
        super(myBertModel, self).__init__()
        self.device = device
        self.num_class = 2
        # bert_config = BertConfig.from_pretrained(bert_path)
        # self.bert = BertModel(bert_config)                #只加载设备，不加载参数

        # self.bert
        # 预训练的神经网络结构 + 权重参数
        # 接收数字张量，输出768维语义特征
        # 看不懂任何文字（中文 / 英文）

        # self.tokenizer
        # 中文词表（文字→数字的映射规则）
        # 把文字转成数字编码（input_ids）、处理长度
        # 无法做语义计算、特征提取

        self.bert = BertModel.from_pretrained(bert_path)    #既BertModel加载设备，又加载参数
        self.tokenizer = BertTokenizer.from_pretrained(bert_path)
        # 是自定义 BERT 分类模型的核心分类头—— 本质是把 BERT 输出的 768 维文本特征，转换成num_class=2维的分类得分（对应负面 / 正面）
        self.out = nn.Sequential(
            nn.Linear(768,num_class)
        )


    # 文本预处理核心方法—— 作用是把原始文本（比如 “这家酒店太差了”）转换成 BERT 模型能识别的张量格式，并放到指定设备（GPU/CPU）上，是 “原始文本” 到 “模型输入” 的必经步骤
    def build_bert_input(self, text):
        # 分词器处理文本，转换成模型可识别的格式
        # return_tensors='pt',指定返回 PyTorch 张量（pt=PyTorch，若写tf则返回 TensorFlow 张量），符合你的模型框架；
        # padding = 'max_length',   核心：不足按max_length填充
        # truncation = True,  配合：超过max_length则截断
        # max_length = 128) , 目标长度：128
        Input = self.tokenizer(text,return_tensors='pt', padding='max_length', truncation=True, max_length=128)
        # 下面几行代码将分词器处理后的 3 个核心张量（input_ids/attention_mask/token_type_ids）从默认的 CPU 设备，迁移到模型运行的目标设备（GPU/CPU），并返回这 3 个对齐设备的张量，供 BERT 模型直接使用
        # 把input_ids张量迁移到指定设备（如GPU）
        input_ids = Input["input_ids"].to(self.device)
        attention_mask =  Input["attention_mask"].to(self.device)
        token_type_ids = Input["token_type_ids"].to(self.device)
        return input_ids, attention_mask, token_type_ids

    def forward(self, text):
        # 文本预处理：将原始文本转换成BERT需要的输入格式
        # 返回值：
        # - input_ids：文本对应的token数字编码
        # - attention_mask：注意力掩码（区分真实token和填充的[PAD]）
        # - token_type_ids：句子类型编码（区分单/双句，单句时全为0）
        input_ids, attention_mask, token_type_ids = self.build_bert_input(text)

        # 送入BERT模型，获取输出
        # return_dict=False：返回元组（而非字典），包含序列特征和池化特征
        # sequence_out：[batch_size, seq_len, hidden_dim] 每个token的特征
        # pooled_output：[batch_size, hidden_dim] 整个句子的聚合特征（<CLS>位置）
        sequence_out, pooled_output = self.bert(input_ids=input_ids,
                                                attention_mask=attention_mask,
                                                token_type_ids=token_type_ids,
                                                return_dict=False)
        # 将句子级的池化特征送入输出层（通常是全连接层），得到最终预测
        out = self.out(pooled_output)
        return out




# def model_Datapara(model, device,  pre_path=None):
#     model = torch.nn.DataParallel(model).to(device)
#     if pre_path != None:
#         model_dict = torch.load(pre_path).module.state_dict()
#         model.module.load_state_dict(model_dict)
#     return model

###train.py###

import torch
import time
import matplotlib.pyplot as plt
import numpy as np
from tqdm import tqdm









def train_val(para):

########################################################
    model = para['model']
    train_loader =para['train_loader']
    val_loader = para['val_loader']
    scheduler = para['scheduler']
    optimizer = para['optimizer']
    loss = para['loss']
    epoch = para['epoch']
    device = para['device']
    save_path = para['save_path']
    max_acc = para['max_acc']
    val_epoch = para['val_epoch']


#################################################
    plt_train_loss = []
    plt_train_acc = []
    plt_val_loss = []
    plt_val_acc = []
    val_rel = []

    for i in range(epoch):
        start_time = time.time()
        model.train()
        train_loss = 0.0
        train_acc = 0.0
        val_acc = 0.0
        val_loss = 0.0
        for batch in tqdm(train_loader):
            model.zero_grad()
            text, labels = batch[0], batch[1].to(device)
            pred = model(text)
            bat_loss = loss(pred, labels)
            bat_loss.backward()
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            train_loss += bat_loss.item()    #.detach 表示去掉梯度
            train_acc += np.sum(np.argmax(pred.cpu().data.numpy(),axis=1)== labels.cpu().numpy())
        plt_train_loss . append(train_loss/train_loader.dataset.__len__())
        plt_train_acc.append(train_acc/train_loader.dataset.__len__())
        if i % val_epoch == 0:
            model.eval()
            with torch.no_grad():
                for batch in tqdm(val_loader):
                    val_text, val_labels = batch[0], batch[1].to(device)
                    val_pred = model(val_text)
                    val_bat_loss = loss(val_pred, val_labels)
                    val_loss += val_bat_loss.cpu().item()

                    val_acc += np.sum(np.argmax(val_pred.cpu().data.numpy(), axis=1) == val_labels.cpu().numpy())
                    val_rel.append(val_pred)

            if val_acc > max_acc:
                torch.save(model, save_path+str(epoch)+"ckpt")
                max_acc = val_acc
            plt_val_loss.append(val_loss/val_loader.dataset.__len__())
            plt_val_acc.append(val_acc/val_loader.dataset.__len__())
            print('[%03d/%03d] %2.2f sec(s) TrainAcc : %3.6f TrainLoss : %3.6f | valAcc: %3.6f valLoss: %3.6f  ' % \
                  (i, epoch, time.time()-start_time, plt_train_acc[-1], plt_train_loss[-1], plt_val_acc[-1], plt_val_loss[-1])
                  )
            # 每训练 50 轮强制保存一次模型（即使不是最优），用于断点续训或回溯训练过程
            if i % 50 == 0:
                torch.save(model, save_path+'-epoch:'+str(i)+ '-%.2f'%plt_val_acc[-1])
        else:
            plt_val_loss.append(plt_val_loss[-1])
            plt_val_acc.append(plt_val_acc[-1])
            print('[%03d/%03d] %2.2f sec(s) TrainAcc : %3.6f TrainLoss : %3.6f   ' % \
                  (i, epoch, time.time()-start_time, plt_train_acc[-1], plt_train_loss[-1])
                  )
    plt.plot(plt_train_loss)
    plt.plot(plt_val_loss)
    plt.title('loss')
    plt.legend(['train', 'val'])
    plt.show()

    plt.plot(plt_train_acc)
    plt.plot(plt_val_acc)
    plt.title('Accuracy')
    plt.legend(['train', 'val'])
    plt.savefig('acc.png')
    plt.show()

###data.py###

from torch.utils.data import Dataset,DataLoader,ConcatDataset
import torch
from sklearn.model_selection import train_test_split
from tqdm import tqdm
import numpy as np



def read_txt_data(path):
    label = []
    data = []
    with open(path, "r", encoding="utf-8") as f:
        for i, line in tqdm(enumerate(f)):
            if i == 0:
                continue  # continue 表示立即执行下一次循环
            if i > 200 and i < 7500:
                continue
            line = line.strip('\n')
            line = line.split(",", 1)  # 1表示分割次数
            label.append(line[0])
            data.append(line[1])
    print(len(label))
    print(len(data))
    return data, label


class JdDataset(Dataset):
    def __init__(self, x, label):
        self.X = x
        label = [int(i) for i in label]
        self.Y = torch.LongTensor(label)

    def __getitem__(self, item):
        return self.X[item], self.Y[item]            # 数据集一般不让返回str， 要写在字典中，或者转为矩阵。

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


def get_dataloader(path, batchsize=1, valSize=0.2):
    x, label = read_txt_data(path)
    train_x, val_x, train_y, val_y = train_test_split(
        x,  # 待拆分的特征数据（你的场景：酒店评论文本列表）
        label,  # 待拆分的标签数据（你的场景：评论的正负标签列表）
        test_size=valSize,  # 验证集占比（你设的valSize=0.2 → 20%数据做验证，80%做训练）
        shuffle=True,  # 拆分前是否打乱数据顺序
        stratify=label  # 按标签“分层抽样”，保证训练/验证集的标签分布和原始数据一致
    )  #valSize为测试集占比例,
    train_set = JdDataset(train_x, train_y)
    val_set = JdDataset(val_x, val_y)
    train_loader = DataLoader(train_set, batch_size=batchsize)
    # PyTorch 中构建训练集批量数据加载器的核心代码，作用是将自定义的训练数据集 train_set（JdDataset 实例）封装成可迭代的 DataLoader 对象，让模型训练时能按指定批次大小（batchsize）逐批读取数据，而非一次性加载全部数据（避免内存溢出）。
    #
    val_loader = DataLoader(val_set, batch_size=batchsize)
    return train_loader, val_loader




if __name__ == '__main__':

    get_dataloader("../jiudian.txt",batchsize=4)