File size: 4,315 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
import os
import sys
import argparse
import csv
import torch
ROOT = os.path.dirname(__file__)
def build_dataset_from_csv(csv_path):
dataset = []
with open(csv_path, newline='', encoding='utf-8') as f:
reader = csv.DictReader(f)
for row in reader:
q = row.get('question') or row.get('Question')
lbl = row.get('label') or row.get('Label')
if q is None or lbl is None:
continue
try:
lbl = int(lbl)
except Exception:
continue
dataset.append({'question': q, 'label': lbl})
return dataset
class EmbedderWrapper:
"""Wrap a SentenceTransformer instance to ensure tensors are on `device`.
The wrapper exposes `encode(texts, convert_to_tensor=True, **kwargs)` same as
the SentenceTransformer API and will move returned torch.Tensor to the device.
"""
def __init__(self, embedder, device):
self.embedder = embedder
self.device = device
def encode(self, texts, convert_to_tensor=True, **kwargs):
emb = self.embedder.encode(texts, convert_to_tensor=convert_to_tensor, **kwargs)
# Move tensor outputs to device (if it's a torch.Tensor)
if convert_to_tensor and hasattr(emb, 'to'):
try:
emb = emb.to(self.device)
except Exception:
pass
return emb
def main():
parser = argparse.ArgumentParser(description='Run SFT training for RouterNetwork')
parser.add_argument('--csv', default='processed_dataset.csv', help='Path to CSV dataset')
parser.add_argument('--model_name', default='all-MiniLM-L6-v2', help='SentenceTransformer model name')
parser.add_argument('--epochs', type=int, default=5)
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument('--checkpoint', default='router_checkpoint.pt', help='Checkpoint output path')
parser.add_argument('--device', default=None, help='Device to use (cpu or cuda). Auto-detect if unset')
args = parser.parse_args()
# Ensure routing_agent directory is importable such that its modules' imports work
routing_agent_path = os.path.join(ROOT, 'routing_agent')
if routing_agent_path not in sys.path:
sys.path.insert(0, routing_agent_path)
# Import modules from routing_agent (sft_trainer expects to import model as a top-level module)
try:
import sft_trainer
import model as model_mod
except Exception as e:
print('Failed importing routing_agent modules:', e)
sys.exit(1)
# Device selection
if args.device:
device = torch.device(args.device)
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Build dataset
csv_path = os.path.join(ROOT, args.csv) if not os.path.isabs(args.csv) else args.csv
if not os.path.exists(csv_path):
print('Dataset CSV not found at', csv_path)
sys.exit(1)
dataset = build_dataset_from_csv(csv_path)
if len(dataset) == 0:
print('No valid rows found in dataset. Exiting.')
sys.exit(1)
# Initialize embedder
try:
from sentence_transformers import SentenceTransformer
except Exception as e:
print('sentence-transformers not installed or failed to import:', e)
sys.exit(1)
embedder = SentenceTransformer(args.model_name)
embedder_wrapper = EmbedderWrapper(embedder, device)
# Infer embedding dim and instantiate model from the same module sft_trainer expects
try:
embed_dim = embedder.get_sentence_embedding_dimension()
except Exception:
# fallback to common dim
embed_dim = 384
model = model_mod.RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)
# Run SFT training
print(f"Starting SFT training on device={device} with model={args.model_name} | epochs={args.epochs} | lr={args.lr}")
sft_trainer.train_sft(model, embedder_wrapper, dataset, epochs=args.epochs, lr=args.lr)
# Save checkpoint
try:
torch.save(model.state_dict(), args.checkpoint)
print('Saved checkpoint to', args.checkpoint)
except Exception as e:
print('Failed to save checkpoint:', e)
if __name__ == '__main__':
main()
|