File size: 1,978 Bytes
975dbdd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | import torch
import torch.optim as optim
import torch.nn as nn
from model import RouterNetwork
def train_sft(model, embedder, dataset, epochs = 5, lr = 1e-3):
"""
dataset: List of dicts [{'question': '...', 'label': 0|1|2}]
"""
device = next(model.parameters()).device
class_weights = torch.tensor([1.0, 1.0, 10.0], device=device, dtype=torch.float)
optimizer = optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss(weight=class_weights)
model.train()
for epoch in range(epochs):
total_loss = 0
correct_predictions = 0
for data in dataset:
# 1. Embed the query (state)
# Using torch.no_grad() for embedder so we only train the router MLP
with torch.no_grad():
query_emb = embedder.encode(data['question'], convert_to_tensor=True)
label = torch.tensor([data['label']], dtype=torch.long).to(query_emb.device)
# 2. Forward pass through the router network
optimizer.zero_grad()
action_logits, _ = model(query_emb) # only care about action_logits (Actor head) for SFT
# Calculate prediction and track accuracy for this epoch
predicted_class = torch.argmax(action_logits, dim=-1)
if predicted_class.item() == label.item():
correct_predictions += 1
# 3. Compute loss and backpropagate
# CrossEntropyLoss expects (N, C) for logits and (N) for labels, so we unsqueeze action_logits to add batch dimension
loss = criterion(action_logits.unsqueeze(0), label)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(dataset)
accuracy = (correct_predictions / len(dataset)) * 100
print(f"SFT Epoch {epoch+1} | Loss: {avg_loss:.4f} | Accuracy: {accuracy:.2f}%") |