import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# 加載預(yù)訓(xùn)練模型和分詞器
model_name = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 準(zhǔn)備目標(biāo)領(lǐng)域數(shù)據(jù)
target_domain_texts = ["This is a medical text.", "Another example from the medical domain."]
target_domain_labels = [1, 0]  # 假設(shè)是二分類任務(wù)

from torch.utils.data import DataLoader, Dataset
from transformers import AdamW

# 自定義數(shù)據(jù)集類
class CustomDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len

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

    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_len,
            return_token_type_ids=False,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt',
        )
        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

# 創(chuàng)建數(shù)據(jù)集和數(shù)據(jù)加載器
max_len = 128
batch_size = 16
train_dataset = CustomDataset(target_domain_texts, target_domain_labels, tokenizer, max_len)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 定義優(yōu)化器
optimizer = AdamW(model.parameters(), lr=2e-5)

# 微調(diào)模型
epochs = 3
for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

from sklearn.metrics import accuracy_score

# 準(zhǔn)備測試數(shù)據(jù)
test_texts = ["This is another medical text.", "More examples for testing."]
test_labels = [1, 0]
test_dataset = CustomDataset(test_texts, test_labels, tokenizer, max_len)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 評估模型
model.eval()
predictions, true_labels = [], []
with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits
        preds = torch.argmax(logits, dim=1)
        predictions.extend(preds.cpu().numpy())
        true_labels.extend(labels.cpu().numpy())

accuracy = accuracy_score(true_labels, predictions)
print(f"Accuracy: {accuracy:.4f}")

# 準(zhǔn)備目標(biāo)任務(wù)數(shù)據(jù)
target_task_texts = ["This is a positive review.", "This is a negative review."]
target_task_labels = [1, 0]  # 假設(shè)是情感分析任務(wù)

# 創(chuàng)建數(shù)據(jù)集和數(shù)據(jù)加載器
train_dataset = CustomDataset(target_task_texts, target_task_labels, tokenizer, max_len)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 微調(diào)模型
epochs = 3
for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        loss.backward()
        optimizer.step()

# 準(zhǔn)備測試數(shù)據(jù)
test_texts = ["This is another positive review.", "This is another negative review."]
test_labels = [1, 0]
test_dataset = CustomDataset(test_texts, test_labels, tokenizer, max_len)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 評估模型
model.eval()
predictions, true_labels = [], []
with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs.logits
        preds = torch.argmax(logits, dim=1)
        predictions.extend(preds.cpu().numpy())
        true_labels.extend(labels.cpu().numpy())

accuracy = accuracy_score(true_labels, predictions)
print(f"Accuracy: {accuracy:.4f}")

from torch.nn import CrossEntropyLoss
from torch.optim import SGD

# 定義對抗訓(xùn)練損失函數(shù)
def adversarial_loss(model, input_ids, attention_mask, labels, epsilon=0.01):
    loss_fn = CrossEntropyLoss()
    outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
    loss = outputs.loss
    loss.backward()
    # 添加對抗擾動
    grad = input_ids.grad
    perturbed_input_ids = input_ids + epsilon * grad.sign()
    perturbed_outputs = model(input_ids=perturbed_input_ids, attention_mask=attention_mask, labels=labels)
    perturbed_loss = perturbed_outputs.loss
    return loss + perturbed_loss

# 使用對抗訓(xùn)練微調(diào)模型
optimizer = SGD(model.parameters(), lr=2e-5)
for epoch in range(epochs):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        loss = adversarial_loss(model, input_ids, attention_mask, labels)
        loss.backward()
        optimizer.step()

from transformers import DistilBertForSequenceClassification

# 加載教師模型和學(xué)生模型
teacher_model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
student_model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=2)

# 定義知識蒸餾損失函數(shù)
def distillation_loss(teacher_logits, student_logits, labels, temperature=2.0, alpha=0.5):
    soft_teacher = torch.softmax(teacher_logits / temperature, dim=-1)
    soft_student = torch.softmax(student_logits / temperature, dim=-1)
    loss_fn = CrossEntropyLoss()
    ce_loss = loss_fn(student_logits, labels)
    kl_loss = torch.nn.functional.kl_div(soft_student.log(), soft_teacher, reduction='batchmean')
    return alpha * ce_loss + (1 - alpha) * kl_loss

# 使用知識蒸餾微調(diào)學(xué)生模型
optimizer = AdamW(student_model.parameters(), lr=2e-5)
for epoch in range(epochs):
    teacher_model.eval()
    student_model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)
        with torch.no_grad():
            teacher_outputs = teacher_model(input_ids=input_ids, attention_mask=attention_mask)
        student_outputs = student_model(input_ids=input_ids, attention_mask=attention_mask)
        loss = distillation_loss(teacher_outputs.logits, student_outputs.logits, labels)
        loss.backward()
        optimizer.step()

# 定義多任務(wù)損失函數(shù)
def multi_task_loss(task1_logits, task2_logits, task1_labels, task2_labels, alpha=0.5):
    loss_fn = CrossEntropyLoss()
    task1_loss = loss_fn(task1_logits, task1_labels)
    task2_loss = loss_fn(task2_logits, task2_labels)
    return alpha * task1_loss + (1 - alpha) * task2_loss

# 使用多任務(wù)學(xué)習(xí)微調(diào)模型
optimizer = AdamW(model.parameters(), lr=2e-5)
for epoch in range(epochs):
    model.train()
    for batch1, batch2 in zip(train_loader1, train_loader2):
        optimizer.zero_grad()
        input_ids1 = batch1['input_ids'].to(device)
        attention_mask1 = batch1['attention_mask'].to(device)
        labels1 = batch1['labels'].to(device)
        input_ids2 = batch2['input_ids'].to(device)
        attention_mask2 = batch2['attention_mask'].to(device)
        labels2 = batch2['labels'].to(device)
        outputs1 = model(input_ids=input_ids1, attention_mask=attention_mask1)
        outputs2 = model(input_ids=input_ids2, attention_mask=attention_mask2)
        loss = multi_task_loss(outputs1.logits, outputs2.logits, labels1, labels2)
        loss.backward()
        optimizer.step()