File size: 2,285 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 | #!/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()
|