Upload routing agent adapters and source code

975dbdd verified 2 months ago

2.29 kB

#!/usr/bin/env python3
import os
import sys
import torch
import pandas as pd

# Ensure routing_agent package is importable
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from ppo_trainer import run_ppo_rollout, load_router_policy
from model import RouterNetwork
from sentence_transformers import SentenceTransformer


def main():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # Prepare embedder and model
    embed_model_name = 'all-MiniLM-L6-v2'
    embedder = SentenceTransformer(embed_model_name)

    embed_dim = embedder.get_sentence_embedding_dimension()
    model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)

    # Optionally load existing checkpoint to continue training
    checkpoint_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint.pt')
    if os.path.exists(checkpoint_path):
        try:
            sd = torch.load(checkpoint_path, map_location='cpu')
            if isinstance(sd, dict) and 'model_state_dict' in sd:
                model.load_state_dict(sd['model_state_dict'], strict=False)
            else:
                model.load_state_dict(sd, strict=False)
            print('Loaded checkpoint', checkpoint_path)
        except Exception as e:
            print('Failed loading checkpoint:', e)

    optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)

    # training_queries = [
    #     "What are the graduation requirements for the 2026 academic year?",
    #     "Show me the syllabus for the Advanced Database course.",
    #     "How do I apply for the on-campus dormitory?"
    # ]
    df = pd.read_csv('df_queries.csv')
    training_queries = df['question'].tolist()
    
    # training_queries = training_queries[:2]  # Limit to first 200 queries for faster testing


    # Run PPO rollouts (this uses simulate_rag_pipeline to call the agent tools)
    run_ppo_rollout(model, embedder, training_queries, optimizer, epochs=40, device=str(device))

    # Save the updated model
    out_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint_pt_after_ppo.pt')
    torch.save({'model_state_dict': model.state_dict()}, out_path)
    print('Saved PPO-updated checkpoint to', out_path)


if __name__ == '__main__':
    main()

	#!/usr/bin/env python3
	import os
	import sys
	import torch
	import pandas as pd

	# Ensure routing_agent package is importable
	ROOT = os.path.dirname(__file__)
	if ROOT not in sys.path:
	sys.path.insert(0, ROOT)

	from ppo_trainer import run_ppo_rollout, load_router_policy
	from model import RouterNetwork
	from sentence_transformers import SentenceTransformer


	def main():
	device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

	# Prepare embedder and model
	embed_model_name = 'all-MiniLM-L6-v2'
	embedder = SentenceTransformer(embed_model_name)

	embed_dim = embedder.get_sentence_embedding_dimension()
	model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(device)

	# Optionally load existing checkpoint to continue training
	checkpoint_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint.pt')
	if os.path.exists(checkpoint_path):
	try:
	sd = torch.load(checkpoint_path, map_location='cpu')
	if isinstance(sd, dict) and 'model_state_dict' in sd:
	model.load_state_dict(sd['model_state_dict'], strict=False)
	else:
	model.load_state_dict(sd, strict=False)
	print('Loaded checkpoint', checkpoint_path)
	except Exception as e:
	print('Failed loading checkpoint:', e)

	optimizer = torch.optim.Adam(model.parameters(), lr=3e-5)

	# training_queries = [
	# "What are the graduation requirements for the 2026 academic year?",
	# "Show me the syllabus for the Advanced Database course.",
	# "How do I apply for the on-campus dormitory?"
	# ]
	df = pd.read_csv('df_queries.csv')
	training_queries = df['question'].tolist()

	# training_queries = training_queries[:2] # Limit to first 200 queries for faster testing


	# Run PPO rollouts (this uses simulate_rag_pipeline to call the agent tools)
	run_ppo_rollout(model, embedder, training_queries, optimizer, epochs=40, device=str(device))

	# Save the updated model
	out_path = os.path.join(os.path.dirname(__file__), 'router_checkpoint_pt_after_ppo.pt')
	torch.save({'model_state_dict': model.state_dict()}, out_path)
	print('Saved PPO-updated checkpoint to', out_path)


	if __name__ == '__main__':
	main()