Upload routing agent adapters and source code

975dbdd verified 2 months ago

4.32 kB

#!/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()

	#!/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()