pip install torch torchvision torchaudio transformers datasets accelerate

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Tokenizer
import numpy as np
from tqdm import tqdm

# ==================== 配置參數(shù) ====================
class Config:
    vocab_size = 50257  # GPT-2 tokenizer 詞匯表大小
    d_model = 768       # 模型維度
    nhead = 12          # 注意力頭數(shù)
    num_layers = 12     # Transformer 層數(shù)
    dropout = 0.1
    batch_size = 8
    seq_len = 128       # 序列長(zhǎng)度
    learning_rate = 3e-4
    num_epochs = 10
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

config = Config()

# ==================== 數(shù)據(jù)集 ====================
class TextDataset(Dataset):
    def __init__(self, texts, tokenizer, max_length=128):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.encodings = []
        
        for text in texts:
            encoding = tokenizer(
                text,
                truncation=True,
                padding='max_length',
                max_length=max_length,
                return_tensors='pt'
            )
            self.encodings.append({
                'input_ids': encoding['input_ids'].squeeze(),
                'attention_mask': encoding['attention_mask'].squeeze()
            })
    
    def __len__(self):
        return len(self.encodings)
    
    def __getitem__(self, idx):
        return self.encodings[idx]

# ==================== Transformer 模型 ====================
class TransformerLM(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # 詞嵌入層
        self.embedding = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_embedding = nn.Embedding(config.seq_len, config.d_model)
        
        # Transformer 編碼器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.d_model,
            nhead=config.nhead,
            dim_feedforward=config.d_model * 4,
            dropout=config.dropout,
            batch_first=True
        )
        self.transformer = nn.TransformerEncoder(encoder_layer, config.num_layers)
        
        # 輸出層
        self.fc_out = nn.Linear(config.d_model, config.vocab_size)
        self.dropout = nn.Dropout(config.dropout)
        
    def forward(self, x, mask=None):
        # 位置編碼
        batch_size, seq_len = x.shape
        positions = torch.arange(0, seq_len, device=x.device).unsqueeze(0)
        
        # 嵌入 + 位置編碼
        x = self.embedding(x) + self.pos_embedding(positions)
        x = self.dropout(x)
        
        # Transformer 編碼
        transformer_out = self.transformer(x, src_key_padding_mask=~mask.bool() if mask is not None else None)
        
        # 輸出預(yù)測(cè)
        output = self.fc_out(transformer_out)
        return output

# ==================== 訓(xùn)練函數(shù) ====================
def train_model(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    
    progress_bar = tqdm(dataloader, desc="Training")
    for batch in progress_bar:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        
        # 創(chuàng)建目標(biāo)（右移一位）
        targets = input_ids[:, 1:].contiguous()
        input_ids = input_ids[:, :-1].contiguous()
        attention_mask = attention_mask[:, :-1].contiguous()
        
        optimizer.zero_grad()
        outputs = model(input_ids, attention_mask)
        
        # 計(jì)算損失（忽略填充位置）
        loss = criterion(outputs.view(-1, config.vocab_size), targets.view(-1))
        loss.backward()
        
        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        total_loss += loss.item()
        progress_bar.set_postfix({'loss': loss.item()})
    
    return total_loss / len(dataloader)

# ==================== 文本生成函數(shù) ====================
def generate_text(model, tokenizer, prompt, max_length=50, temperature=1.0, top_k=50):
    model.eval()
    with torch.no_grad():
        # 編碼輸入提示
        input_ids = tokenizer.encode(prompt, return_tensors='pt').to(config.device)
        generated = input_ids
        
        for _ in range(max_length):
            # 獲取模型輸出
            outputs = model(generated)
            next_token_logits = outputs[:, -1, :] / temperature
            
            # Top-k 采樣
            if top_k > 0:
                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
                next_token_logits[indices_to_remove] = -float('Inf')
            
            # Softmax + 采樣
            probs = torch.softmax(next_token_logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            
            # 檢查是否生成結(jié)束符
            if next_token.item() == tokenizer.eos_token_id:
                break
                
            generated = torch.cat([generated, next_token], dim=-1)
        
        return tokenizer.decode(generated[0], skip_special_tokens=True)

# ==================== 主訓(xùn)練流程 ====================
def main():
    # 初始化 tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
    tokenizer.pad_token = tokenizer.eos_token
    
    # 準(zhǔn)備示例數(shù)據(jù)（實(shí)際使用時(shí)替換為真實(shí)數(shù)據(jù)集）
    sample_texts = [
        "Artificial intelligence is transforming the world.",
        "Machine learning models require large amounts of data.",
        "Natural language processing enables computers to understand human language.",
        "Deep learning has achieved remarkable success in various domains.",
        "Transformer architecture revolutionized sequence modeling."
    ] * 100  # 重復(fù)以創(chuàng)建足夠數(shù)據(jù)
    
    # 創(chuàng)建數(shù)據(jù)集和數(shù)據(jù)加載器
    dataset = TextDataset(sample_texts, tokenizer, config.seq_len)
    dataloader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True)
    
    # 初始化模型
    model = TransformerLM(config).to(config.device)
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)
    optimizer = optim.AdamW(model.parameters(), lr=config.learning_rate)
    
    # 訓(xùn)練循環(huán)
    print(f"Starting training on {config.device}...")
    for epoch in range(config.num_epochs):
        avg_loss = train_model(model, dataloader, optimizer, criterion, config.device)
        print(f"Epoch {epoch+1}/{config.num_epochs}, Average Loss: {avg_loss:.4f}")
        
        # 每 2 個(gè) epoch 生成示例文本
        if (epoch + 1) % 2 == 0:
            prompt = "Artificial intelligence"
            generated_text = generate_text(model, tokenizer, prompt, max_length=30)
            print(f"Generated text: {generated_text}\n")
    
    # 保存模型
    torch.save(model.state_dict(), 'transformer_lm.pth')
    print("Model saved successfully!")

if __name__ == "__main__":
    main()

# 絕對(duì)位置編碼 vs 相對(duì)位置編碼
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
    
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

# 多頭注意力計(jì)算過程
def scaled_dot_product_attention(q, k, v, mask=None):
    """
    q, k, v: [batch_size, seq_len, d_k]
    """
    d_k = q.size(-1)
    scores = torch.matmul(q, k.transpose(-2, -1)) / np.sqrt(d_k)  # [B, L, L]
    
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    attn_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attn_weights, v)
    return output, attn_weights

# 防止梯度爆炸
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

# 預(yù)熱 + 余弦退火
def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
    def lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * (current_step - num_warmup_steps) / 
                                               float(max(1, num_training_steps - num_warmup_steps)))))
    return optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

# 忽略填充 token 的損失計(jì)算
criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

# 控制生成多樣性
next_token_logits = outputs[:, -1, :] / temperature
# temperature < 1: 更確定性
# temperature > 1: 更隨機(jī)性

# Top-k 采樣
def top_k_sampling(logits, k):
    indices_to_remove = logits < torch.topk(logits, k)[0][..., -1, None]
    logits[indices_to_remove] = -float('Inf')
    return logits

# Top-p (Nucleus) 采樣
def top_p_sampling(logits, p):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = 0
    indices_to_remove = sorted_indices_to_remove.scatter(
        dim=-1, index=sorted_indices, src=sorted_indices_to_remove
    )
    logits[indices_to_remove] = -float('Inf')
    return logits

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        outputs = model(input_ids)
        loss = criterion(outputs.view(-1, vocab_size), targets.view(-1))
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

# 多 GPU 訓(xùn)練
model = nn.DataParallel(model)
# 或使用 DistributedDataParallel (更高效)
model = nn.parallel.DistributedDataParallel(model)

# 動(dòng)態(tài)量化
quantized_model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments

# 加載預(yù)訓(xùn)練模型
model = GPT2LMHeadModel.from_pretrained('gpt2')
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token

# 訓(xùn)練參數(shù)
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
)

# 訓(xùn)練器
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

trainer.train()

from transformers import pipeline

# 使用 pipeline 進(jìn)行文本生成
generator = pipeline('text-generation', model='gpt2', device=0)

result = generator(
    "Artificial intelligence is",
    max_length=50,
    num_return_sequences=1,
    temperature=0.7,
    top_k=50,
    top_p=0.95
)
print(result[0]['generated_text'])