Upload routing agent adapters and source code

975dbdd verified 2 months ago

1.98 kB

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}%")

	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}%")