"""Evaluation and metrics pipeline for CAT V3.""" from __future__ import annotations import time import torch from typing import Dict, List, Any import numpy as np from cat_v3.dataset import DOMAINS, ConceptVocabulary, SimpleCharTokenizer def compute_routing_metrics(preds: torch.Tensor, targets: torch.Tensor) -> Dict[str, float]: """Computes F1-score, Precision, and Recall for multi-label expert classification.""" preds_bool = preds.bool().cpu().numpy() targets_bool = targets.bool().cpu().numpy() tp = np.logical_and(preds_bool, targets_bool).sum() fp = np.logical_and(preds_bool, np.logical_not(targets_bool)).sum() fn = np.logical_and(np.logical_not(preds_bool), targets_bool).sum() precision = tp / max(tp + fp, 1e-9) recall = tp / max(tp + fn, 1e-9) f1 = 2 * (precision * recall) / max(precision + recall, 1e-9) exact_match = np.all(preds_bool == targets_bool, axis=-1).mean() return { "router_precision": float(precision), "router_recall": float(recall), "router_f1": float(f1), "router_exact_match": float(exact_match) } def compute_path_metrics(pred_paths: List[List[str]], target_paths: List[List[str]]) -> Dict[str, float]: """Computes F1, precision, recall, and exact sequence match for reasoning paths.""" tp_c, fp_c, fn_c = 0, 0, 0 exact_matches = 0 for pred, target in zip(pred_paths, target_paths): p_set = set(pred) t_set = set(target) tp = len(p_set.intersection(t_set)) fp = len(p_set - t_set) fn = len(t_set - p_set) tp_c += tp fp_c += fp fn_c += fn if pred == target: exact_matches += 1 precision = tp_c / max(tp_c + fp_c, 1e-9) recall = tp_c / max(tp_c + fn_c, 1e-9) f1 = 2 * (precision * recall) / max(precision + recall, 1e-9) exact_match = exact_matches / max(len(target_paths), 1) return { "path_precision": float(precision), "path_recall": float(recall), "path_f1": float(f1), "path_exact_match": float(exact_match) } def evaluate_model( model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, vocab: ConceptVocabulary, tokenizer: SimpleCharTokenizer, ) -> Dict[str, float]: """Evaluates router routing, reasoning path, and decoder text generation metrics.""" model.eval() device = next(model.parameters()).device all_router_preds = [] all_router_targets = [] all_pred_paths = [] all_target_paths = [] exact_response_matches = 0 total_samples = 0 latencies = [] with torch.no_grad(): for batch in dataloader: input_ids = batch["input_ids"].to(device) attention_mask = batch["attention_mask"].to(device) router_target = batch["router_target"].to(device) path_ids = batch["path_ids"].to(device) response_ids = batch["response_ids"].to(device) # Measure latency start_time = time.perf_counter() outputs = model.generate_response( input_ids=input_ids, attention_mask=attention_mask, router_top_k=2, router_threshold=0.5 ) latency = (time.perf_counter() - start_time) * 1000.0 latencies.append(latency / input_ids.size(0)) # Router evaluations all_router_preds.append(outputs["router_mask"]) all_router_targets.append(router_target) # Decode paths & responses for b in range(input_ids.size(0)): total_samples += 1 # Active experts target active_domains = [DOMAINS[idx] for idx, val in enumerate(router_target[b].tolist()) if val > 0] # Decoded target path t_ids = path_ids[b].cpu().tolist() t_path = vocab.decode_path(t_ids) all_target_paths.append(t_path) # Decode predicted path from active experts or fusion layer fused_ids = outputs["fused_concept_ids"][b].cpu().tolist() f_path = vocab.decode_path(fused_ids) all_pred_paths.append(f_path) # Evaluate decoder exact response match target_response = tokenizer.decode(response_ids[b].cpu().tolist()) generated_response = tokenizer.decode(outputs["generated_tokens"][b].cpu().tolist()) if target_response.strip().lower() == generated_response.strip().lower(): exact_response_matches += 1 router_preds_tensor = torch.cat(all_router_preds, dim=0) router_targets_tensor = torch.cat(all_router_targets, dim=0) router_metrics = compute_routing_metrics(router_preds_tensor, router_targets_tensor) path_metrics = compute_path_metrics(all_pred_paths, all_target_paths) eval_metrics = { **router_metrics, **path_metrics, "decoder_exact_match": exact_response_matches / max(total_samples, 1), "inference_latency_ms": float(np.mean(latencies)), } return eval_metrics def run_single_inference( model: torch.nn.Module, question: str, vocab: ConceptVocabulary, tokenizer: SimpleCharTokenizer, ) -> Dict[str, Any]: """Runs a single inference query and returns the detailed multi-expert reasoning report.""" model.eval() device = next(model.parameters()).device input_ids, attention_mask = tokenizer.encode(question, max_length=32) input_ids = input_ids.unsqueeze(0).to(device) attention_mask = attention_mask.unsqueeze(0).to(device) with torch.no_grad(): outputs = model.generate_response( input_ids=input_ids, attention_mask=attention_mask, router_top_k=2, router_threshold=0.5 ) # 1. Router probabilities router_probs = outputs["router_probs"][0].cpu().tolist() active_mask = outputs["router_mask"][0].cpu().tolist() activated_domains = [DOMAINS[idx] for idx, val in enumerate(active_mask) if val > 0] domain_probs = {DOMAINS[idx]: router_probs[idx] for idx in range(len(DOMAINS))} # 2. Expert local paths expert_paths = {} for idx, domain in enumerate(DOMAINS): if active_mask[idx]: path_ids = outputs["expert_reports"][domain]["predicted_path"][0].cpu().tolist() path = vocab.decode_path(path_ids) expert_paths[domain] = path # 3. Concept fusion report fusion_report = model.fusion.get_symbolic_report( vocab=vocab, expert_reports=outputs["expert_reports"], router_mask=outputs["router_mask"], domain_names=DOMAINS )[0] # 4. Final Text Generation gen_tokens = outputs["generated_tokens"][0].cpu().tolist() response = tokenizer.decode(gen_tokens) return { "question": question, "activated_domains": activated_domains, "domain_probabilities": domain_probs, "expert_paths": expert_paths, "fusion_report": fusion_report, "answer": response }