from transformers import AutoModelForSequenceClassification
from transformers import AutoTokenizer

model_path = 'D:/Models/chinese-roberta-wwm-ext'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path, num_labels=15)  # 分类个数

import re
from sklearn.utils import shuffle
import pandas as pd

label_dic = {
    'news_story':0,
    'news_culture':1,
    'news_entertainment':2,
    'news_sports':3,
    'news_finance':4,
    'news_house':5,
    'news_car':6,
    'news_edu':7,
    'news_tech':8,
    'news_military':9,
    'news_travel':10,
    'news_world':11,
    'stock':12,
    'news_agriculture':13,
    'news_game':14
}

def get_train_data(file_path, col_num):
    content = []
    label = []
    with open(file_path, "r", encoding="utf-8") as f:
        num = 0
        for i in f.readlines():
            if num > col_num:
                break
            lines = i.split("_!_")
            content.append(re.sub('[^\u4e00-\u9fa5]',"",lines[3]))  # 去除非中文
            label.append(label_dic.get(lines[2]))
            num += 1
    return content,label
        
content,label = get_train_data("./file/toutiao_cat_data.txt", 8000)
data = pd.DataFrame({"content":content,"label":label})
data = shuffle(data)

train_data = tokenizer(data.content.to_list(), padding = "max_length", max_length = 100, truncation=True ,return_tensors = "pt")
train_label = data.label.to_list()

import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler, SequentialSampler

batch_size = 16
train = TensorDataset(train_data["input_ids"], train_data["attention_mask"], torch.tensor(train_label))
train_sampler = RandomSampler(train)
train_dataloader = DataLoader(train, sampler=train_sampler, batch_size=batch_size)

# 定义优化器
from torch.optim import AdamW
optimizer = AdamW(model.parameters(), lr=1e-4)
# 定义学习率和训练轮数
num_epochs = 1
from transformers import get_scheduler
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    name="linear", optimizer=optimizer, num_warmup_steps=0, num_training_steps=num_training_steps
)

from tqdm import tqdm

device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
model.to(device)
for epoch in range(num_epochs):
    total_loss = 0
    model.train()
    with tqdm(list(enumerate(train_dataloader)),ncols=100) as _tqdm:
        for step, batch in _tqdm:
            _tqdm.set_description('epoch {}/{}'.format(epoch+1, num_epochs))
            if not step == 0:
                cur_loss = total_loss/(step*batch_size)
                avg_train_loss = total_loss / len(train_dataloader)
                _tqdm.set_postfix(loss=cur_loss, avg_loss=avg_train_loss)
                _tqdm.update(1)
            else:
                _tqdm.set_postfix(loss=0.00000)
            b_input_ids = batch[0].to(device)
            b_input_mask = batch[1].to(device)
            b_labels = batch[2].to(device)
            model.zero_grad()        
            outputs = model(b_input_ids, 
                        token_type_ids=None, 
                        attention_mask=b_input_mask, 
                        labels=b_labels)

            loss = outputs.loss       
            total_loss += loss.item()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            lr_scheduler.step()

import numpy as np

test = tokenizer("通过考研大纲谈农学复习之化学部分年农学将第二次实行全国统一考试这使农学考生的复习备考十分迷茫",return_tensors="pt",padding="max_length",max_length=100)
test.to(device)

model.eval()
with torch.no_grad():  
    outputs = model(test["input_ids"], 
                    token_type_ids=None, 
                    attention_mask=test["attention_mask"])

logits = outputs["logits"].cpu()
pred_flat = np.argmax(logits,axis=1).numpy().squeeze()
print(pred_flat.tolist())
print(list(label_dic.keys())[list(label_dic.values()).index(pred_flat.tolist())])