import torch
import torch.nn as nn

class GPTClassifier(nn.Module):
    """在 GPT 基础上加分类头"""
    def __init__(self, base_model, num_classes=2, dropout=0.1):
        super().__init__()
        self.base_model = base_model
        
        # 冻结基座模型 (不训练 GPT 部分)
        for param in self.base_model.parameters():
            param.requires_grad = False
        
        # 分类头
        emb_dim = base_model.config["emb_dim"]
        self.classifier = nn.Sequential(
            nn.LayerNorm(emb_dim),
            nn.Linear(emb_dim, num_classes),
        )
    
    def forward(self, x):
        # x: (batch, seq_len) token IDs
        with torch.no_grad():
            # 取最后一个 token 的输出 (像 [CLS] 一样)
            base_out = self.base_model(x)  # (batch, seq, emb_dim)
        last_token = base_out[:, -1, :]  # (batch, emb_dim)
        return self.classifier(last_token)

print("✓ GPTClassifier 定义完成")
print("实际使用时: classifier = GPTClassifier(gpt_model, num_classes=2)")

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

# 模拟数据集
data = [
    ("Free money!!! Click here now", 1),  # 垃圾
    ("Hey, how are you?", 0),             # 正常
    ("Congratulations! You won a prize", 1),
    ("Meeting at 3pm tomorrow", 0),
    ("Buy cheap watches now!!!", 1),
    ("Can we reschedule the call?", 0),
]

class SpamDataset(Dataset):
    def __init__(self, data, max_length=64):
        self.enc = tiktoken.get_encoding("gpt2")
        self.max_length = max_length
        self.inputs = []
        self.labels = []
        
        for text, label in data:
            ids = self.enc.encode(text)
            # 填充或截断到固定长度
            padded = ids[:max_length] + [0] * (max_length - len(ids))
            self.inputs.append(padded)
            self.labels.append(label)
    
    def __len__(self): return len(self.inputs)
    def __getitem__(self, idx):
        return torch.tensor(self.inputs[idx]), torch.tensor(self.labels[idx])

dataset = SpamDataset(data)
loader = DataLoader(dataset, batch_size=2, shuffle=True)

x, y = next(iter(loader))
print(f"批次: input {x.shape}, labels {y}")

def train_classifier(model, loader, n_epochs=10, lr=1e-3, device='cpu'):
    model.to(device)
    model.train()
    optimizer = torch.optim.AdamW(model.classifier.parameters(), lr=lr)
    loss_fn = nn.CrossEntropyLoss()
    
    for epoch in range(n_epochs):
        total_loss = 0.0
        correct = 0
        total = 0
        
        for x, y in loader:
            x, y = x.to(device), y.to(device)
            
            optimizer.zero_grad()
            logits = model(x)
            loss = loss_fn(logits, y)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            _, predicted = torch.max(logits, 1)
            correct += (predicted == y).sum().item()
            total += y.size(0)
        
        avg_loss = total_loss / len(loader)
        acc = correct / total
        print(f"Epoch {epoch+1:2d}/{n_epochs} | Loss: {avg_loss:.4f} | Acc: {acc:.2%}")
    return model

print("\n训练函数定义完成。在实际 notebook 中执行:")
print("model = GPTClassifier(gpt_model, num_classes=2)")
print("train_classifier(model, loader, n_epochs=10, device='mps')")

def predict_spam(model, text, device='cpu'):
    """预测一条短信是否为垃圾"""
    enc = tiktoken.get_encoding("gpt2")
    ids = enc.encode(text)[:64]
    padded = ids + [0] * (64 - len(ids))
    x = torch.tensor([padded]).to(device)
    
    model.eval()
    with torch.no_grad():
        logits = model(x)
        prob = torch.softmax(logits, dim=-1)
        pred = torch.argmax(logits, dim=-1).item()
    
    return {
        "text": text,
        "prediction": "垃圾短信" if pred == 1 else "正常",
        "confidence": prob[0][pred].item()
    }

# 演示 (需要已训练的模型)
print("预测函数定义完成。使用示例:")
print('predict_spam(model, "Free money!!!")')

def save_model(model, path="models/spam_classifier.pt"):
    """保存模型权重"""
    torch.save({
        'classifier_state_dict': model.classifier.state_dict(),
        'config': model.base_model.config,
    }, path)
    print(f"模型已保存到 {path}")

def load_model(base_model, path="models/spam_classifier.pt"):
    """加载模型权重"""
    model = GPTClassifier(base_model)
    checkpoint = torch.load(path, map_location='cpu')
    model.classifier.load_state_dict(checkpoint['classifier_state_dict'])
    print(f"模型已从 {path} 加载")
    return model

print("✓ 保存/加载函数定义完成")

# %%ai openai-chat-custom:Qwen3.5-9B-Q4_K_M.gguf
# 解释为什么微调后的模型需要重新做安全审计？基座模型和新加的分类头可能带来什么新的安全风险？