| """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) |
| |
| |
| 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)) |
| |
| |
| all_router_preds.append(outputs["router_mask"]) |
| all_router_targets.append(router_target) |
| |
| |
| for b in range(input_ids.size(0)): |
| total_samples += 1 |
| |
| |
| active_domains = [DOMAINS[idx] for idx, val in enumerate(router_target[b].tolist()) if val > 0] |
| |
| |
| t_ids = path_ids[b].cpu().tolist() |
| t_path = vocab.decode_path(t_ids) |
| all_target_paths.append(t_path) |
| |
| |
| fused_ids = outputs["fused_concept_ids"][b].cpu().tolist() |
| f_path = vocab.decode_path(fused_ids) |
| all_pred_paths.append(f_path) |
| |
| |
| 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 |
| ) |
| |
| |
| 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))} |
| |
| |
| 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 |
| |
| |
| fusion_report = model.fusion.get_symbolic_report( |
| vocab=vocab, |
| expert_reports=outputs["expert_reports"], |
| router_mask=outputs["router_mask"], |
| domain_names=DOMAINS |
| )[0] |
| |
| |
| 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 |
| } |
|
|
