S0L009/Luna-GNN-Scorer-InferenceAPI
0
1"""2Helpful model training, testing, and inference functions.3"""4 5import torch6import torch.nn as nn7from torch_geometric.data import Data8 9import numpy as np10import pickle as pkl11from typing import Callable, Optional12 13from data import RegenerativeData, Neo4jAggregator14from model import DirWeighedGNN, MultiheadScoring, flexible_device15from loss import masked_mse, unmasked_mse16 17 18def inference_step(19 score_model: nn.Module,20 embeddings: torch.Tensor,21 edge_index: torch.Tensor,22 edge_weights: torch.Tensor,23 categorical_src_indices: list[torch.Tensor],24 categorical_trg_indices: list[torch.Tensor],25 gnn_model: Optional[nn.Module] = None,26 all_pairs: bool = False,27 training: bool = True28) -> list[torch.Tensor]:29 """30 Performs an inference step to compute predicted scores using a scoring model.31 32 Parameters33 ----------34 score_model : nn.Module35 The scoring model used to compute scores based on embeddings.36 embeddings : torch.Tensor37 The input vertex embeddings.38 edge_index : torch.Tensor39 Tensor containing edge indices that define the graph structure40 in COO format.41 edge_weights : torch.Tensor42 Weights associated with each edge in the graph.43 categorical_src_indices : list[torch.Tensor]44 Tensor containing source vertex indices for each category.45 categorical_trg_indices : list[torch.Tensor]46 Tensor containing target vertex indices for each category.47 gnn_model : nn.Module, optional48 The GNN model to extract refinined vertex embeddings, by default None.49 all_pairs : bool, optional50 Flag to indicate whether to compute scores for all pairs, by default False.51 training : bool, optional52 Flag to indicate if the model is in training mode, by default True.53 54 Returns55 -------56 list[torch.Tensor]57 The list of predicted scores for each category.58 """59 60 if gnn_model is not None:61 embeddings = gnn_model(62 x=embeddings,63 edge_index=edge_index,64 edge_weights=edge_weights,65 training=training66 )67 68 pred_scores = score_model(69 embeddings=embeddings,70 categorical_src_indices=categorical_src_indices,71 categorical_trg_indices=categorical_trg_indices,72 all_pairs=all_pairs,73 training=training74 )75 76 return pred_scores77 78 79def loss_step(80 data: RegenerativeData,81 score_model: nn.Module,82 criterion: Callable,83 gnn_model: Optional[nn.Module] = None,84 training: bool = True85) -> torch.Tensor:86 """87 Computes the model prediction loss on input data.88 89 Parameters90 ----------91 data: RegenerativeData92 A `RegenerativeData` object generated by a `GraphAggregator`93 when called with `generate_graph_data`.94 score_model: nn.Module95 The scoring model based on vertex embeddings.96 criterion: Callable97 The loss function.98 gnn_model: nn.Module, optional99 The GNN model to extract vertex embeddings. If not provided,100 the `score_model` is expected to use graph embeddings.101 Defaults to None.102 training: bool, optional103 Whether or not the loss is computed for model training.104 Defaults to True.105 106 Returns107 -------108 loss: torch.Tensor109 The loss for the input data.110 """111 pred_scores = inference_step(112 score_model=score_model,113 embeddings=data.x,114 edge_index=data.edge_index,115 edge_weights=data.masked_edge_attr,116 categorical_src_indices=data.sources,117 categorical_trg_indices=data.targets,118 gnn_model=gnn_model,119 all_pairs=False,120 training=training121 )122 123 losses = [criterion(pred, y, mask)124 for pred, y, mask125 in zip(pred_scores, data.y, data.mask)]126 loss = sum(losses) / len(losses)127 128 return loss129 130 131def train_scores_predictor(132 data: RegenerativeData,133 score_model: nn.Module,134 optim: torch.optim.Optimizer,135 criterion: Callable,136 steps: int,137 gnn_model: nn.Module = None,138 val_data: list[RegenerativeData] = [],139 val_criteria: list[Callable] = [],140 val_names: list[str] = [],141 early_stop: Callable = None,142 print_every: int = None,143 verbose: bool = True,144 seed: int = 142937145) -> nn.Module | tuple[nn.Module, nn.Module]:146 """147 Train the scoring predictor model.148 149 Parameters150 ----------151 data: RegenerativeData152 The training data.153 score_model: nn.Module154 The scoring model used to turn embeddings into scores.155 optim: torch.optim.Optimizer156 The optimizer to use for training.157 criterion: Callable158 The loss function.159 steps: int160 The number of training steps.161 gnn_model: nn.Module, optional162 The GNN model used to extract embeddings, by default None.163 val_data: list[RegenerativeData], optional164 List of validation datasets, by default [].165 val_criteria: list[Callable], optional166 The validation loss functions, by default []. If [],167 will use `criterion` for all validation datasets.168 val_names: list[str], optional169 The names of the validation datasets, by default [].170 early_stop: Callable, optional171 The early stopping criterion. If None, the training will run to completion.172 The default is None.173 174 The criterion should accept a dictionary object and return True175 if early stopping should be triggered. The dictionary contains the values of176 various metrics throughout the training process. Specifically, let ``name``177 be the name of any metric name (where ``name`` is either 'loss' or a name in178 ``val_names``), then the dictionary will contain the following keys:179 180 * ``name``: A (step, value) tuple of the current step and the current value181 of the metric.182 183 * ``name_low``: A (step, value) tuple of the lowest value of the metric and184 the step.185 186 * ``name_high``: A (step, value) tuple of the highest value of the metric187 and the step.188 189 * ``name_hist``: A list of all values of the metric, indexed by step.190 191 Example::192 193 def early_stop(metrics):194 # stop if the `val` metric has increased from its195 # lowest value by more than 0.1196 return metrics['val'][1] - metrics['val_low'][1] > 0.1197 198 print_every: int, optional199 How often to print the loss, by default None. If None,200 will print every 10% of the steps.201 verbose: bool, optional202 Whether to print messages during training, by default True.203 seed: int, optional204 The seed to use for random number generation, by default 142937.205 206 Returns207 -------208 float or tuple[float, float]209 The training loss, and optionally the validation loss if provided.210 """ 211 torch.manual_seed(seed)212 213 if print_every is None:214 print_every = max(1, steps // 10)215 216 if not val_criteria:217 val_criteria = [criterion for _ in range(len(val_data))]218 219 if not val_names:220 val_names = [f'Val-{i}' for i in range(len(val_data))]221 222 if early_stop is None:223 early_stop = lambda x: False224 225 metrics = {}226 for name in ['loss'] + val_names:227 metrics = metrics | {228 f'{name}': (-1, None),229 f'{name}_low': (-1, float('inf')),230 f'{name}_high': (-1, -float('inf')),231 f'{name}_hist': [],232 }233 234 def update_metric(name: str, step: int, value: float):235 metrics[f'{name}'] = (step, value)236 if metrics[f'{name}_low'][1] > value:237 metrics[f'{name}_low'] = (step, value)238 if metrics[f'{name}_high'][1] < value:239 metrics[f'{name}_high'] = (step, value)240 metrics[f'{name}_hist'].append(value)241 242 for step in range(steps):243 data.regenerate()244 optim.zero_grad()245 loss = loss_step(246 data=data,247 gnn_model=gnn_model,248 score_model=score_model,249 criterion=criterion,250 training=True251 )252 loss.backward()253 optim.step()254 255 if verbose and step % print_every == 0:256 print(f"[{step:4d}] Loss: {loss:.5f}", end=' ')257 258 loss = loss.detach().cpu().item()259 update_metric('loss', step, loss)260 261 with torch.no_grad(): 262 for nm, vd, cr in zip(val_names, val_data, val_criteria):263 vd.regenerate()264 val_loss = loss_step(265 data=vd,266 gnn_model=gnn_model,267 score_model=score_model,268 criterion=cr,269 training=False270 )271 val_loss = val_loss.detach().cpu().item()272 update_metric(nm, step, val_loss)273 274 if verbose and step % print_every == 0:275 print(f"{nm}: {val_loss:.5f}", end=' ')276 277 if verbose and step % print_every == 0:278 print()279 280 if step > 0 and early_stop(metrics):281 if verbose:282 print('Stopping early...')283 break284 285 return metrics286 287 288def predict_scores(289 data: RegenerativeData|Data,290 score_model: nn.Module,291 src_indices: torch.Tensor = None,292 trg_indices: torch.Tensor = None,293 gnn_model: nn.Module = None,294 all_pairs: bool = True295) -> list[torch.Tensor]:296 """297 Compute scores for the given data using.298 299 Parameters300 ----------301 data: RegenerativeData|Data302 The data to compute scores for.303 score_model: nn.Module304 The scoring model used to turn embeddings into scores.305 src_indices: torch.Tensor, optional306 The source vertex indices for each category. Defaults to None307 and reads from `data.sources`.308 trg_indices: torch.Tensor, optional309 The target vertex indices for each category. Defaults to None310 and reads from `data.targets`.311 gnn_model: nn.Module, optional312 The GNN model used to extract embeddings, by default None.313 all_pairs: bool, optional314 Whether to compute scores for all pairs of source and target315 vertices, by default True.316 317 Returns318 -------319 scores: list[torch.Tensor]320 The predicted scores.321 """322 if src_indices is None:323 src_indices = data.sources324 if trg_indices is None:325 trg_indices = data.targets326 327 with torch.no_grad():328 scores = inference_step(329 score_model=score_model,330 embeddings=data.x,331 edge_index=data.edge_index,332 edge_weights=data.edge_attr,333 categorical_src_indices=src_indices,334 categorical_trg_indices=trg_indices,335 gnn_model=gnn_model,336 all_pairs=all_pairs,337 training=False338 )339 340 return scores341 342def get_query_matrices(343 data: RegenerativeData|Data,344 score_model: nn.Module,345 src_indices: torch.Tensor = None,346 trg_indices: torch.Tensor = None,347 gnn_model: nn.Module = None,348 all_pairs: bool = True349) -> list[torch.Tensor]:350 """351 Compute scores for the given data using.352 353 Parameters354 ----------355 data: RegenerativeData|Data356 The data to compute scores for.357 score_model: nn.Module358 The scoring model used to turn embeddings into scores.359 src_indices: torch.Tensor, optional360 The source vertex indices for each category. Defaults to None361 and reads from `data.sources`.362 trg_indices: torch.Tensor, optional363 The target vertex indices for each category. Defaults to None364 and reads from `data.targets`.365 gnn_model: nn.Module, optional366 The GNN model used to extract embeddings, by default None.367 all_pairs: bool, optional368 Whether to compute scores for all pairs of source and target369 vertices, by default True.370 371 Returns372 -------373 scores: list[torch.Tensor]374 The predicted scores.375 """376 if src_indices is None:377 src_indices = data.sources378 if trg_indices is None:379 trg_indices = data.targets380 381 with torch.no_grad():382 if gnn_model is not None:383 embeddings = gnn_model(384 x=data.x,385 edge_index=data.edge_index,386 edge_weights=data.edge_attr,387 training=False388 )389 390 queries = score_model(391 embeddings=embeddings,392 categorical_src_indices=src_indices, # user393 categorical_trg_indices=trg_indices, # categories394 all_pairs=all_pairs,395 training=False,396 just_query=True397 )398 399 obj = {400 'queries': queries,401 'embeddings': embeddings402 }403 torch.save(obj, 'query_cache.pt')404 405 return obj406 407def neo4j_training_pipeline(408 db_path: str,409 graph_steps: int = 100,410 raw_steps: int = 100411):412 agg = Neo4jAggregator(db_path)413 agg.collect()414 device = flexible_device()415 416 masked_data = agg.generate_graph_data(417 device=device,418 prob_masked=0.6,419 )420 421 unmasked_data = agg.generate_graph_data(422 device=device,423 prob_masked=0.0424 )425 426 DEMB = agg.vertex_attrs[0].shape[0]427 DMODEL = 512428 429 gnn_model = DirWeighedGNN(430 in_channels=DEMB,431 hidden_channels=1024,432 out_channels=DEMB433 ).to(device)434 435 score_model = MultiheadScoring(436 nheads=agg.category_count,437 demb=DEMB,438 dmodel=DMODEL439 ).to(device)440 441 raw_score_model = MultiheadScoring(442 nheads=agg.category_count,443 demb=DEMB,444 dmodel=DMODEL445 ).to(device)446 447 optim_graph = torch.optim.Adam(448 list(gnn_model.parameters()) + list(score_model.parameters()),449 lr=1e-3450 )451 452 train_scores_predictor(453 data=masked_data,454 score_model=score_model,455 gnn_model=gnn_model,456 optim=optim_graph,457 criterion=masked_mse,458 steps=100459 )460 461 optim_raw = torch.optim.Adam(462 raw_score_model.parameters(),463 lr=1e-3464 )465 466 train_scores_predictor(467 data=masked_data,468 score_model=score_model,469 gnn_model=gnn_model,470 optim=optim_graph,471 criterion=masked_mse,472 steps=graph_steps473 )474 475 train_scores_predictor(476 data=unmasked_data,477 score_model=raw_score_model,478 optim=optim_raw,479 criterion=unmasked_mse,480 steps=raw_steps481 )482 483 torch.save(484 {485 'gnn': gnn_model.state_dict(),486 'score_model': score_model.state_dict(),487 'raw_score_model': raw_score_model.state_dict(),488 },489 'models.pt'490 )491 492 torch.save(493 {494 'optim_graph': optim_graph.state_dict(),495 'optim_raw': optim_raw.state_dict()496 },497 'optims.pt'498 )499 500def get_cached_scores(501 cache_path: str,502 user_id: str,503 category: str,504 target_ids: list[str],505):506 return np.random.random()507 508 with open(cache_path, 'rb') as f:509 pass510 obj = pkl.load(f)511 512 def score(user_id, target_id):513 user = obj.id2vertex[user_id]514 target = obj.id2vertex[target_id]515 category = obj.id2category[category]516 return obj.scores[user][category][target]517 518 return {519 trg_id: score(user_id, trg_id)520 for trg_id in target_ids521 }522 523 524def recommend_users(525 venue_id: str,526 user_id: str,527 num_users: int = 10528):529 ...530