from datasets import load_dataset  
 
 from transformers import AutoTokenizer, AutoModelForSequenceClassification  
 from transformers import DistilBertForSequenceClassification, DistilBertConfig  
 
 import torch  
 import torch.nn as nn  
 import torch.optim as optim  
 from torch.utils.data import DataLoader  
 
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support

data = load_dataset("shawhin/phishing-site-classification")

# 使用Nvidia GPU  
 device = torch.device('cuda')  
 
 # 加載教師模型和分詞器
 model_path = "shawhin/bert-phishing-classifier_teacher"  
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)  
 teacher_model = AutoModelForSequenceClassification.from_pretrained(model_path)  
                                                  .to(device)

# 加載學(xué)生模型
 my_config = DistilBertConfig(n_heads=8, n_layers=4) # 每層減少4個(gè)頭,總共減少2層  
 
 student_model = DistilBertForSequenceClassification  
                                    .from_pretrained("distilbert-base-uncased",  
                                     config=my_config,)  
                                    .to(device)

# 定義文本預(yù)處理函數(shù)
 def preprocess_function(examples):  
     return tokenizer(examples["text"], padding='max_length', truncation=True)  
 
 # 對所有數(shù)據(jù)集進(jìn)行標(biāo)記化
 tokenized_data = data.map(preprocess_function, batched=True)  
 tokenized_data.set_format(type='torch',  
                           columns=['input_ids', 'attention_mask', 'labels'])

# 評估模型性能的函數(shù)
 def evaluate_model(model, dataloader, device):  
     model.eval()  # 將模型設(shè)置為評估模式
     all_preds = []  
     all_labels = []  
 
     # 禁用梯度計(jì)算
     with torch.no_grad():  
         for batch in dataloader:  
             input_ids = batch['input_ids'].to(device)  
             attention_mask = batch['attention_mask'].to(device)  
             labels = batch['labels'].to(device)  
 
             # 前向傳播獲取logits  
             outputs = model(input_ids, attention_mask=attention_mask)  
             logits = outputs.logits  
 
             # 獲取預(yù)測結(jié)果
             preds = torch.argmax(logits, dim=1).cpu().numpy()  
             all_preds.extend(preds)  
             all_labels.extend(labels.cpu().numpy())  
 
     # 計(jì)算評估指標(biāo)
     accuracy = accuracy_score(all_labels, all_preds)  
     precision, recall, f1, _ = precision_recall_fscore_support(all_labels,  
                                                               all_preds,  
                                                 average='binary')  
 
     return accuracy, precision, recall, f1

# 計(jì)算蒸餾損失和硬標(biāo)簽損失的函數(shù)
 def distillation_loss(student_logits, teacher_logits,  
                       true_labels, temperature, alpha):  
     # 從教師logits計(jì)算軟目標(biāo)
     soft_targets = nn.functional.softmax(teacher_logits / temperature, dim=1)  
     student_soft = nn.functional.log_softmax(student_logits / temperature, dim=1)  
 
     # 蒸餾的KL散度損失
     distill_loss = nn.functional.kl_div(student_soft,  
                                     soft_targets,  
                                     reduction='batchmean') * (temperature ** 2)  
 
     # 硬標(biāo)簽的交叉熵?fù)p失
     hard_loss = nn.CrossEntropyLoss()(student_logits, true_labels)  
 
     # 結(jié)合損失
     loss = alpha * distill_loss + (1.0 - alpha) * hard_loss  
 
     return loss

# 超參數(shù)
 batch_size = 32  
 lr = 1e-4  
 num_epochs = 5  
 temperature = 2.0  
 alpha = 0.5  
 
 # 定義優(yōu)化器
 optimizer = optim.Adam(student_model.parameters(), lr=lr)  
 
 # 創(chuàng)建訓(xùn)練數(shù)據(jù)加載器
 dataloader = DataLoader(tokenized_data['train'], batch_size=batch_size)  
 # 創(chuàng)建測試數(shù)據(jù)加載器
 test_dataloader = DataLoader(tokenized_data['test'], batch_size=batch_size)

# 將學(xué)生模型設(shè)置為訓(xùn)練模式
 student_model.train()  
 
 # 訓(xùn)練模型
 for epoch in range(num_epochs):  
     for batch in dataloader:  
         # 準(zhǔn)備輸入
         input_ids = batch['input_ids'].to(device)  
         attention_mask = batch['attention_mask'].to(device)  
         labels = batch['labels'].to(device)  
 
         # 禁用教師模型的梯度計(jì)算
         with torch.no_grad():  
             teacher_outputs = teacher_model(input_ids,  
                                             attention_mask=attention_mask)  
             teacher_logits = teacher_outputs.logits  
 
         # 學(xué)生模型前向傳播
         student_outputs = student_model(input_ids,  
                                         attention_mask=attention_mask)  
         student_logits = student_outputs.logits  
 
         # 計(jì)算蒸餾損失
         loss = distillation_loss(student_logits, teacher_logits, labels,  
                                   temperature, alpha)  
 
         # 反向傳播
         optimizer.zero_grad()  
         loss.backward()  
         optimizer.step()  
 
     print(f"第 {epoch + 1} 輪訓(xùn)練完成,損失: {loss.item()}")  
 
     # 評估教師模型
     teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =  
                          evaluate_model(teacher_model, test_dataloader, device)  
 
     print(f"教師模型 (測試集) - 準(zhǔn)確率: {teacher_accuracy:.4f},  
                                 精確率: {teacher_precision:.4f},  
                                 召回率: {teacher_recall:.4f},  
                                 F1分?jǐn)?shù): {teacher_f1:.4f}")  
 
     # 評估學(xué)生模型
     student_accuracy, student_precision, student_recall, student_f1 =  
                          evaluate_model(student_model, test_dataloader, device)  
       
     print(f"學(xué)生模型 (測試集) - 準(zhǔn)確率: {student_accuracy:.4f},  
                                 精確率: {student_precision:.4f},  
                                 召回率: {student_recall:.4f},  
                                 F1分?jǐn)?shù): {student_f1:.4f}")  
     print("\n")  
 
     # 將學(xué)生模型重新設(shè)置為訓(xùn)練模式
     student_model.train()

# 創(chuàng)建驗(yàn)證數(shù)據(jù)加載器
 validation_dataloader = DataLoader(tokenized_data['validation'], batch_size=8)  
 
 # 評估教師模型
 teacher_accuracy, teacher_precision, teacher_recall, teacher_f1 =  
                    evaluate_model(teacher_model, validation_dataloader, device)  
 print(f"教師模型 (驗(yàn)證集) - 準(zhǔn)確率: {teacher_accuracy:.4f},  
                            精確率: {teacher_precision:.4f},  
                            召回率: {teacher_recall:.4f},  
                            F1分?jǐn)?shù): {teacher_f1:.4f}")  
 
 # 評估學(xué)生模型
 student_accuracy, student_precision, student_recall, student_f1 =  
                    evaluate_model(student_model, validation_dataloader, device)  
 print(f"學(xué)生模型 (驗(yàn)證集) - 準(zhǔn)確率: {student_accuracy:.4f},  
                            精確率: {student_precision:.4f},  
                            召回率: {student_recall:.4f},  
                            F1分?jǐn)?shù): {student_f1:.4f}")

from transformers import BitsAndBytesConfig  
 
 # 以4位精度加載模型
 nf4_config = BitsAndBytesConfig(  
     load_in_4bit=True,  
     bnb_4bit_quant_type="nf4",  
     bnb_4bit_compute_dtype = torch.bfloat16,  
     bnb_4bit_use_double_quant=True  
 )  
 
 model_nf4 = AutoModelForSequenceClassification.from_pretrained(model_id,  
                                                 device_map=device,  
                                                 quantization_config=nf4_config)

# 評估量化后的學(xué)生模型
 quantized_accuracy, quantized_precision, quantized_recall, quantized_f1 =  
                        evaluate_model(model_nf4, validation_dataloader, device)  
 
 print("量化后性能")  
 print(f"準(zhǔn)確率: {quantized_accuracy:.4f},  
         精確率: {quantized_precision:.4f},  
         召回率: {quantized_recall:.4f},  
         F1分?jǐn)?shù): {quantized_f1:.4f}")