routing-agent / run_inference.py

Upload routing agent adapters and source code

975dbdd verified 2 months ago

4.48 kB

import os
import sys
import time
import torch
from sentence_transformers import SentenceTransformer
import pandas as pd

# ── 1. Path setup (same pattern as run_ppo-2.py) ────────────────────────────
ROOT = os.path.dirname(__file__)
if ROOT not in sys.path:
    sys.path.insert(0, ROOT)

from model import RouterNetwork
from grpo_trainer import simulate_rag_pipeline   # reuse your existing tool caller

# ── 2. Config ────────────────────────────────────────────────────────────────
CHECKPOINT  = os.path.join(ROOT, "adapters/router_grpo_epoch_16.pt")   # ← change epoch here
EMBED_MODEL = "all-MiniLM-L6-v2"
DEVICE      = "cuda" if torch.cuda.is_available() else "cpu"

ACTION_NAMES = {
    0: "Guidelines  (academic rules / visa / scholarships)",
    1: "Web-Search   (news / announcements)",
    2: "CBNU Database (courses / departments / contacts)",
}

# ── 3. Load embedder + model ─────────────────────────────────────────────────
embedder  = SentenceTransformer(EMBED_MODEL)
embed_dim = embedder.get_sentence_embedding_dimension()
model     = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(DEVICE)

# Load checkpoint (supports both raw state_dict and wrapped dict)
sd = torch.load(CHECKPOINT, 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)

model.eval()
print(f"✅  Loaded checkpoint: {CHECKPOINT}")
print(f"    Device: {DEVICE}\n")

# ── 4. Single-query inference function ───────────────────────────────────────
def route_query(query: str, call_tool: bool = True):
    """
    Route one query through the loaded GRPO checkpoint.
    Args:
        query     : The user question string.
        call_tool : If True, also calls simulate_rag_pipeline and prints the answer.
                    Set False for routing-only / fast testing.
    Returns:
        action_idx (int), action_name (str), context (str), answer (str)
    """
    # Embed
    with torch.no_grad():
        emb = embedder.encode(query, convert_to_tensor=True).to(DEVICE)
        action, log_prob, _ = model.get_action(emb)
        action_idx = int(action.item())

    action_name = ACTION_NAMES[action_idx]
    print(f"Query   : {query}")
    print(f"→ Routed to Action {action_idx}: {action_name}")
    print(f"  log_prob = {log_prob.item():.4f}")

    context, answer = "", ""
    if call_tool:
        print("  Calling tool …")
        context, answer = simulate_rag_pipeline(query, action_idx)
        print(f"  Answer  : {answer[:300]}")

    return action_idx, action_name, context, answer


# ── 5. Main: test with one or more queries ───────────────────────────────────
if __name__ == "__main__":
    test_queries = [
        "What are the graduation requirements for international students?",
        "How do I register for courses next semester?",
        "What is the latest news from CBNU?",
    ]
    test_queries = [
        "2학기 개강일은?",
        "알바하고 싶은데 어떤 조건이 필요한가?",
        "아르바이트 하려는데 조건은?",
        "시간제 취업 방법은?",
        "건강보험 어떻게 가입해?",
        "성적확인 기간은?"
    ]
    
    eval_df = pd.read_csv('ko_eval_queries.csv')
    test_queries = eval_df['question'].tolist()
    results = []
    for idx, row in eval_df.iterrows():
        q = row.get('question')
        ground = row.get('action_label')
        start = time.perf_counter()
        pred_idx, pred_name, context, answer = route_query(q, call_tool=False)
        end = time.perf_counter()
        elapsed = end - start
        results.append({
            'question': q,
            'ground_truth': ground,
            'predicted_label': pred_idx,
            'time_took': elapsed,
        })
        
        
    out_df = pd.DataFrame(results)
    out_df.to_csv('ko_eval_with_timing_results_epoch_16.csv', index=False)
    print(f"Wrote results to eval_with_timing_results.csv ({len(out_df)} rows)")

	import os
	import sys
	import time
	import torch
	from sentence_transformers import SentenceTransformer
	import pandas as pd

	# ── 1. Path setup (same pattern as run_ppo-2.py) ────────────────────────────
	ROOT = os.path.dirname(__file__)
	if ROOT not in sys.path:
	sys.path.insert(0, ROOT)

	from model import RouterNetwork
	from grpo_trainer import simulate_rag_pipeline # reuse your existing tool caller

	# ── 2. Config ────────────────────────────────────────────────────────────────
	CHECKPOINT = os.path.join(ROOT, "adapters/router_grpo_epoch_16.pt") # ← change epoch here
	EMBED_MODEL = "all-MiniLM-L6-v2"
	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

	ACTION_NAMES = {
	0: "Guidelines (academic rules / visa / scholarships)",
	1: "Web-Search (news / announcements)",
	2: "CBNU Database (courses / departments / contacts)",
	}

	# ── 3. Load embedder + model ─────────────────────────────────────────────────
	embedder = SentenceTransformer(EMBED_MODEL)
	embed_dim = embedder.get_sentence_embedding_dimension()
	model = RouterNetwork(embedding_dim=embed_dim, num_classes=3).to(DEVICE)

	# Load checkpoint (supports both raw state_dict and wrapped dict)
	sd = torch.load(CHECKPOINT, 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)

	model.eval()
	print(f"✅ Loaded checkpoint: {CHECKPOINT}")
	print(f" Device: {DEVICE}\n")

	# ── 4. Single-query inference function ───────────────────────────────────────
	def route_query(query: str, call_tool: bool = True):
	"""
	Route one query through the loaded GRPO checkpoint.

	Args:
	query : The user question string.
	call_tool : If True, also calls simulate_rag_pipeline and prints the answer.
	Set False for routing-only / fast testing.

	Returns:
	action_idx (int), action_name (str), context (str), answer (str)
	"""
	# Embed
	with torch.no_grad():
	emb = embedder.encode(query, convert_to_tensor=True).to(DEVICE)
	action, log_prob, _ = model.get_action(emb)
	action_idx = int(action.item())

	action_name = ACTION_NAMES[action_idx]
	print(f"Query : {query}")
	print(f"→ Routed to Action {action_idx}: {action_name}")
	print(f" log_prob = {log_prob.item():.4f}")

	context, answer = "", ""
	if call_tool:
	print(" Calling tool …")
	context, answer = simulate_rag_pipeline(query, action_idx)
	print(f" Answer : {answer[:300]}")

	return action_idx, action_name, context, answer


	# ── 5. Main: test with one or more queries ───────────────────────────────────
	if __name__ == "__main__":
	test_queries = [
	"What are the graduation requirements for international students?",
	"How do I register for courses next semester?",
	"What is the latest news from CBNU?",
	]
	test_queries = [
	"2학기 개강일은?",
	"알바하고 싶은데 어떤 조건이 필요한가?",
	"아르바이트 하려는데 조건은?",
	"시간제 취업 방법은?",
	"건강보험 어떻게 가입해?",
	"성적확인 기간은?"
	]

	eval_df = pd.read_csv('ko_eval_queries.csv')
	test_queries = eval_df['question'].tolist()
	results = []
	for idx, row in eval_df.iterrows():
	q = row.get('question')
	ground = row.get('action_label')
	start = time.perf_counter()
	pred_idx, pred_name, context, answer = route_query(q, call_tool=False)
	end = time.perf_counter()
	elapsed = end - start
	results.append({
	'question': q,
	'ground_truth': ground,
	'predicted_label': pred_idx,
	'time_took': elapsed,
	})


	out_df = pd.DataFrame(results)
	out_df.to_csv('ko_eval_with_timing_results_epoch_16.csv', index=False)
	print(f"Wrote results to eval_with_timing_results.csv ({len(out_df)} rows)")