S0L009/Luna-GNN-Scorer-InferenceAPI
0
1"""2Data collection utilities.3"""4 5 6import torch7from torch import Tensor8from torch_geometric.data import Data9from sentence_transformers import SentenceTransformer10 11import random12import numpy as np13import warnings14import requests15from tqdm import tqdm16from dataclasses import dataclass, field17from typing import Any, Optional, Hashable, Callable18 19from model import flexible_device20from model import MultiheadScoring21from db_utils import generic_doc_ref22 23 24def create_masked_data(25 data: Tensor,26 device,27 val_masked: int = 0,28 prob_masked: int = 029):30 """31 Create a masked version of the given data tensor.32 33 Parameters34 ----------35 data: Tensor36 The input data tensor.37 device : torch.device38 The device to place the tensors on.39 val_masked : int, optional40 The value to set the masked elements to, by default 0.41 prob_masked : int, optional42 The probability of an element being masked, by default 0.43 44 Returns45 -------46 masked_data: Tensor47 The masked data.48 mask: Tensor49 The mask, where `True` indicates the element is masked.50 """51 if len(data) == 0:52 return (data.to(device),53 torch.tensor([], dtype=torch.bool).to(device))54 55 mask = torch.bernoulli(torch.ones_like(data) * prob_masked) == 156 data[mask] = val_masked57 return data.to(device), mask.to(device)58 59 60class RegenerativeData(Data):61 """62 A data structure that extends PyTorch Geometric's Data class,63 used for storing graph data and generating masked edge attributes.64 65 Attributes66 ----------67 device : torch.device68 The device on which the data is stored.69 x : torch.Tensor70 The vertex embeddings.71 edge_index : torch.Tensor72 The graph connectivity in COO format.73 edge_attr : torch.Tensor74 The original edge attributes.75 sources : torch.Tensor76 The source vertex indices.77 targets : torch.Tensor78 The target vertex indices.79 y : torch.Tensor80 The edge labels.81 prob_masked : float, optional82 The probability of masking an edge attribute.83 val_masked : float, optional84 The value to set the masked elements to.85 edge_splitter : Callable, optional86 The function to split the edge masks into a mask for each category.87 mask : torch.Tensor88 The mask where `True` indicates the element is masked.89 masked_edge_attr : torch.Tensor90 The edge attributes with masked values.91 """92 device: torch.device93 x: torch.Tensor94 edge_index: torch.Tensor95 edge_attr: torch.Tensor96 sources: torch.Tensor97 targets: torch.Tensor98 y: torch.Tensor99 prob_masked: float100 val_masked: float101 edge_splitter: Callable102 mask: torch.Tensor103 masked_edge_attr: torch.Tensor104 105 def __init__(self, *args, **kwargs):106 super().__init__(*args, **kwargs)107 108 def regenerate(self):109 """110 Randomly regenerates `masked_edge_attr` and `mask` according111 to the current `prob_masked` and `val_masked` parameters.112 All other attributes remain unchanged.113 """114 self.masked_edge_attr, self.mask = create_masked_data(115 data=self.edge_attr.clone().to(self.device),116 val_masked=self.val_masked,117 prob_masked=self.prob_masked,118 device=self.device119 )120 self.mask = self.edge_splitter(self.mask)121 122 123@dataclass124class GraphAggregator():125 """126 A graph data structure that allows for the incremental addition of127 vertices and edges together with their associated data and categories.128 The graph is stored in a way that allows for efficient generation of129 PyTorch Geometric data objects, with the option of masking out random130 edge weights and separating targets by category.131 132 Attributes133 ----------134 id2category : dict135 Maps category IDs to internal category representation.136 category2id : list137 Maps internal category representations to category ID.138 category_count : int139 Counter for categories.140 vertex2category : list141 Maps vertices to their categories.142 category_selector : list143 Selects edges by category. The category of an edge is determined by144 the category of its target vertex. The i-th element of this list145 contains all the edges in the i-th category.146 id2vertex : dict147 Maps vertex IDs to internal representations.148 vertex2id : list149 Maps internal vertex representations to vertex ID.150 vertex_count : int151 Counter for vertices.152 vertex_attrs : list153 Stores attributes of vertices.154 id2edge : dict155 Maps edge keys to internal representations.156 edge2id : list157 Maps internal edge representations to edge ID.158 edge_count : int159 Counter for edges.160 edge_attrs : list161 Stores attributes of edges.162 sources : list163 Source vertices for edges.164 targets : list165 Target vertices for edges.166 167 Example::168 169 graph = GraphAggregator()170 graph.add_vertex('Nico', category_id='user', data=[1, 2, 3])171 graph.add_vertex('VITAL Climbing', category_id='recreation', data=[4, 5, 6])172 graph.add_vertex('Pink Taco', category_id='restaurant', data=[7, 8, 9])173 graph.add_vertex('Vincent', category_id='user', data=[10, 11, 12])174 graph.add_edge('Vincent', 'VITAL Climbing', data=0.8)175 graph.add_edge('Nico', 'VITAL Climbing', data=0.6)176 graph.add_edge('Vincent', 'Pink Taco', data=0.7)177 graph.add_edge('Nico', 'Vincent', data=1)178 data = graph.generate_graph_data(device='cpu', prob_masked=0.5)179 """180 torch.rand181 id2category: dict = field(default_factory=dict)182 category2id: list = field(default_factory=list)183 category_count: int = 0184 vertex2category: list = field(default_factory=list)185 category_selector: list = field(default_factory=list)186 187 id2vertex: dict = field(default_factory=dict)188 vertex2id: list = field(default_factory=list)189 vertex_count: int = 0190 vertex_attrs: list = field(default_factory=list)191 192 id2edge: dict = field(default_factory=dict)193 edge2id: list = field(default_factory=list)194 edge_count: int = 0195 edge_attrs: list = field(default_factory=list)196 sources: list = field(default_factory=list)197 targets: list = field(default_factory=list)198 199 def exists_vertex(200 self,201 vertex_id: Hashable202 ) -> bool:203 """204 Checks if a vertex with the given id exists in the graph.205 206 Parameters207 ----------208 vertex_id : Hashable209 The id of the vertex to check.210 211 Returns212 -------213 bool214 True if a vertex with the given id exists, False otherwise.215 """216 return vertex_id in self.id2vertex217 218 def add_vertex(219 self,220 vertex_id: Hashable,221 category_id: Optional[Hashable] = None,222 data: Optional[Any] = None223 ):224 """225 Add a vertex to the graph.226 227 Parameters228 ----------229 vertex_id : Hashable230 The id of the vertex.231 category_id : Optional[Hashable], optional232 The id of the category of the vertex, by default None233 data : Any, optional234 The attributes of the vertex, by default None235 236 Raises237 ------238 ValueError239 If the vertex already exists.240 """241 if vertex_id in self.id2vertex:242 warnings.warn(f"\nVertex {vertex_id} already exists, skipping "243 "any actions. To modify the data of an existing "244 "vertex explicitly, use `update_vertex`.")245 246 else:247 if category_id is None:248 raise ValueError(249 f"A new vertex ({vertex_id}) must have a category id."250 )251 252 if category_id not in self.id2category:253 category = self.category_count254 self.category_count += 1255 self.id2category[category_id] = category256 self.category2id.append(category_id)257 self.category_selector.append([])258 259 vertex = self.vertex_count260 self.vertex_count += 1261 self.id2vertex[vertex_id] = vertex262 self.vertex2id.append(vertex_id)263 self.vertex_attrs.append(data)264 265 category = self.id2category[category_id]266 self.vertex2category.append(category)267 268 def update_vertex(269 self,270 vertex_id: Hashable,271 data: Any272 ):273 """274 Update a vertex's attributes.275 276 Parameters277 ----------278 vertex_id : Hashable279 The id of the vertex.280 data : Any281 The new attributes of the vertex.282 """283 self.vertex_attrs[self.id2vertex[vertex_id]] = data284 285 def _edge_key_scheme(286 self,287 src_id: Hashable,288 trg_id: Hashable289 ) -> Hashable:290 """291 The key scheme used to store edges in the `id2edge` dictionary.292 293 Parameters294 ----------295 src_id : Hashable296 The id of the source vertex.297 trg_id : Hashable298 The id of the target vertex.299 300 Returns301 -------302 Hashable303 The key used to store the edge in the `id2edge` dictionary.304 """305 return (src_id, trg_id)306 307 def exists_edge(308 self,309 src_id: Hashable,310 trg_id: Hashable311 ) -> bool:312 """313 Check if an edge exists between two vertices.314 315 Parameters316 ----------317 src_id : Hashable318 The id of the source vertex.319 trg_id : Hashable320 The id of the target vertex.321 322 Returns323 -------324 bool325 True if the edge exists, False otherwise.326 """327 return self._edge_key_scheme(src_id, trg_id) in self.id2edge328 329 def add_edge(330 self,331 src_id: Hashable,332 trg_id: Hashable,333 data: Optional[Any] = None334 ):335 """336 Add an edge to the graph.337 338 Parameters339 ----------340 src_id : Hashable341 The id of the source vertex.342 trg_id : Hashable343 The id of the target vertex.344 data : Optional[Any], optional345 The attributes of the edge, by default None346 347 Raises348 ------349 ValueError350 If the edge already exists.351 """352 _key = self._edge_key_scheme(src_id, trg_id)353 354 if _key in self.id2edge:355 warnings.warn(f"\nEdge {_key} already exists, skipping any "356 "actions. To modify an existing edge explicitly, "357 "use `update_edge` instead.")358 359 else:360 edge = self.edge_count361 self.edge_count += 1362 self.id2edge[_key] = edge363 self.edge2id.append(_key)364 self.edge_attrs.append(data)365 366 src_vertex = self.id2vertex[src_id]367 trg_vertex = self.id2vertex[trg_id]368 self.sources.append(src_vertex)369 self.targets.append(trg_vertex)370 371 trg_cat = self.vertex2category[trg_vertex]372 self.category_selector[trg_cat].append(edge)373 374 def update_edge(375 self,376 src_id: Hashable,377 trg_id: Hashable,378 data: Optional[Any]379 ) -> None:380 """381 Update an edge's attributes.382 383 The edge is identified by its source and target vertex ids.384 If the edge does not exist, a `ValueError` is raised.385 386 Parameters387 ----------388 src_id : Hashable389 Source vertex id.390 trg_id : Hashable391 Target vertex id.392 data : Optional[Any]393 New edge attributes.394 395 Returns396 -------397 None398 """399 _key = self._edge_key_scheme(src_id, trg_id)400 if _key not in self.id2edge:401 raise ValueError(f"Edge {_key} does not exist.")402 self.edge_attrs[self.id2edge[_key]] = data403 404 def get_vertex_data(self, vertex_id):405 """406 Get the attributes of a vertex.407 408 Parameters409 ----------410 vertex_id : Hashable411 The id of the vertex.412 413 Returns414 -------415 Tensor416 The attributes of the vertex.417 418 Raises419 ------420 ValueError421 If the vertex does not exist.422 """423 if vertex_id not in self.id2vertex:424 raise ValueError(f"Vertex {vertex_id} does not exist.")425 return self.vertex_attrs[self.id2vertex[vertex_id]]426 427 def get_edge_data(self, src_id, trg_id):428 """429 Get the attributes of an edge.430 431 Parameters432 ----------433 src_id : Hashable434 The id of the source vertex.435 trg_id : Hashable436 The id of the target vertex.437 438 Returns439 -------440 Tensor441 The attributes of the edge.442 443 Raises444 ------445 ValueError446 If the edge does not exist.447 """448 _key = self._edge_key_scheme(src_id, trg_id)449 if _key not in self.id2edge:450 raise ValueError(f"Edge {_key} does not exist.")451 return self.edge_attrs[self.id2edge[_key]]452 453 def make_edge_splitter(454 self,455 split_categories: bool = True456 ):457 """458 Creates a function that splits a given tensor of edge attributes into459 different categories. The Tensor must have the same length as the number460 of edges in the graph.461 462 Parameters463 ----------464 split_categories : bool, optional465 Whether to split the tensor into different categories, by default True466 467 Returns468 -------469 Callable[[Tensor], List[Tensor]]470 A function that takes a tensor and splits it into different categories471 """472 if split_categories:473 return lambda x : [474 x[selector]475 for selector in self.category_selector476 ]477 else:478 return lambda x : [x]479 480 def make_edge_joiner(481 self,482 split_categories: bool = True483 ):484 """485 Creates a function that joins a given tensor of edge attributes split by486 different categories. The final Tensor will have the same length as the number487 of edges in the graph.488 489 Parameters490 ----------491 split_categories : bool, optional492 Whether the tensor has been split into different categories, by default True493 494 Returns495 -------496 Callable[[List[Tensor]], Tensor]497 A function that takes a list of tensors and joins them into a single tensor498 """499 if split_categories:500 def _join(x: Tensor):501 joined = torch.zeros_like(torch.cat(x, dim=0))502 for cat, selector in enumerate(self.category_selector):503 joined[selector] = x[cat]504 return joined505 return _join506 else:507 return lambda x : x[0]508 509 def generate_graph_data(510 self,511 device,512 prob_masked: float = 0.0,513 val_masked: float = 0.0,514 split_categories: bool = True515 ) -> RegenerativeData:516 """517 Generate a Regenerative PyTorch Geometric data object from the graph.518 519 Parameters520 ----------521 device : torch.device522 The device to which the data should be moved.523 prob_masked : float, optional524 Probability of masking an edge attribute, by default 0.0.525 val_masked : float, optional526 Value to set the masked elements to, by default 0.0.527 split_categories : bool, optional528 Whether to split y into separate tensors based on categories,529 by default True.530 531 Returns532 -------533 RegenerativeData534 The Regenerative PyTorch Geometric data object.535 """ 536 x = torch.tensor(np.array(self.vertex_attrs), dtype=torch.float32).to(device)537 edge_index = torch.tensor([self.sources, self.targets], dtype=torch.long).to(device)538 edge_attr = torch.tensor(np.array(self.edge_attrs), dtype=torch.float32).to(device)539 540 splitter = self.make_edge_splitter(split_categories=split_categories)541 542 data = RegenerativeData(543 x=x,544 edge_index=edge_index,545 edge_attr=edge_attr,546 prob_masked=prob_masked,547 val_masked=val_masked,548 edge_splitter=splitter,549 device=device550 )551 552 data.regenerate()553 554 data.sources = splitter(torch.tensor(self.sources, dtype=torch.long).to(device))555 data.targets = splitter(torch.tensor(self.targets, dtype=torch.long).to(device))556 data.y = splitter(edge_attr.clone().to(device))557 558 return data559 560 def collect(self):561 """562 Call this method to collect data.563 564 To be overloaded in subclasses.565 """566 pass567 568 569# Functions that require a separate instance per graph570GraphAggregatorContextualFns = (571 'exists_vertex',572 'add_vertex',573 'update_vertex',574 'exists_edge',575 'add_edge',576 'update_edge',577 'get_vertex_data',578 'get_edge_data',579 'make_edge_splitter',580 'generate_graph_data'581)582 583def make_multigraph(584 Aggregator: type,585 tracked_fns: tuple[str] = GraphAggregatorContextualFns586):587 """588 Creates a new class that wraps around the Aggregator class589 and provides a unified interface for multiple graphs. Tracked590 functions will expose a new `graph_id` keyword argument which591 can be used to control specified graphs.592 593 594 Parameters595 ----------596 Aggregator : type597 The class to wrap around.598 tracked_fns : tuple[str], optional599 The functions to track and pass through the select() method,600 by default GraphAggregatorContextualFns.601 602 Returns603 -------604 _MultiAggregator605 The new class.606 607 This wrapper can be used as both a decorator or a function.608 609 Example::610 611 @make_multigraph612 class CustomAggregator(GraphAggregator):613 ...614 615 # OR #616 617 MultiCustomAggregator = make_multigraph(CustomAggregator)618 """619 620 class _MultiAggregator(Aggregator):621 def __init__(622 self,623 *args,624 graph_ids: tuple = None,625 **kwargs626 ):627 """628 Initialize a multi-graph aggregator object.629 630 Parameters631 ----------632 *args : tuple633 Additional positional arguments to pass to the Aggregator constructor.634 graph_ids : tuple, optional635 A tuple of graph identifiers used as keys to access each graph, by default636 None (thereby instantiating just the single, main graph).637 **kwargs : dict638 Additional keyword arguments to pass to the Aggregator constructor.639 """640 641 self.graph_ids = graph_ids642 643 super().__init__(*args, **kwargs)644 self.aggregators = {645 id: Aggregator(*args, **kwargs)646 for id in graph_ids647 }648 649 def select(650 self,651 graph_id: str|None = None652 ) -> GraphAggregator:653 """654 Select a specific graph by its identifier.655 656 Parameters657 ----------658 graph_id : str, optional659 The identifier of the graph to select, by default None (thereby selecting660 the main graph).661 662 Returns663 -------664 GraphAggregator665 The selected graph. If graph_id is None, returns self.666 """667 if graph_id is None:668 return self669 if graph_id not in self.aggregators:670 raise ValueError(f"Graph {graph_id} does not exist.")671 return self.aggregators[graph_id]672 673 def __getattribute__(self, attr):674 """675 Overloads the __getattribute__ method to intercept method calls676 for methods listed in tracked_fns. When such a method is called,677 this method intercepts the call and wraps the original method678 with a new method that accepts a graph_id argument.679 680 See the docstring for _wrapped_fn below for more details.681 """682 if attr in tracked_fns:683 _orig_fn = super().__getattribute__(attr)684 _aggregators = self.aggregators685 686 def _access_aggregator(687 graph_id: str688 ):689 if graph_id not in _aggregators:690 raise ValueError(f"Graph {graph_id} does not exist.")691 return _aggregators[graph_id]692 693 def _expand_single(694 obj: list|tuple|str695 ) -> list:696 if isinstance(obj, list) or isinstance(obj, tuple):697 return [_access_aggregator(graph_id=id)698 for id in obj]699 return [_access_aggregator(graph_id=obj)]700 701 def _collapse_single(702 obj: list703 ):704 if len(obj) == 1:705 return obj[0]706 return obj707 708 def wrapped_fn(709 *args,710 graph_id: list|tuple|str|None = None,711 **kwargs712 ):713 """714 A wrapped version of the original function that accepts a715 graph_id argument, which can be either a single string,716 a list of strings, or a tuple of strings.717 718 If graph_id is None, the original function is called on719 the main graph.720 721 If graph_id is a list or tuple, the function is called722 once for each graph_id in the list, and the results are723 combined into a list of results.724 725 If graph_id is a string, the function is called once with726 that graph_id, and the result is returned directly.727 """728 if graph_id is None:729 return _orig_fn(*args, **kwargs)730 731 return _collapse_single([732 agg.__getattribute__(attr)(*args, **kwargs)733 for agg in _expand_single(graph_id)734 ])735 return wrapped_fn736 737 return super().__getattribute__(attr)738 739 return _MultiAggregator740 741def make_multigraph_with(742 tracked_fns: tuple[str] = GraphAggregatorContextualFns743):744 """745 Returns a function that takes a graph aggregator class and returns a new746 multigraph aggregator class that is identical to the original, but with747 the methods listed in `tracked_fns` overriden to accept a `graph_id`748 argument. This can be used to instantiate a decorator which extends custom749 aggregators to multiple graphs.750 751 Parameters752 ----------753 tracked_fns : tuple[str]754 A tuple of method names to override, by default755 GraphAggregatorContextualFns.756 757 Returns758 -------759 Callable[[type], type]760 A function that takes a graph aggregator class and returns a new761 multigraph aggregator class.762 """763 return lambda Aggregator: make_multigraph(Aggregator, tracked_fns)764 765 766MultiGraphAggregator = make_multigraph(GraphAggregator)767 768 769class ListTracker():770 """771 A utility class for tracking and aggregating values in a list.772 773 Attributes774 ----------775 list : list776 The list to store the values in.777 778 Methods779 -------780 append(val)781 Appends a value to the list.782 sum()783 Returns the sum of all the values in the list.784 min()785 Returns the minimum value in the list.786 max()787 Returns the maximum value in the list.788 count_none()789 Returns the number of None values in the list.790 count_not_none()791 Returns the number of non-None values in the list.792 mean()793 Returns the mean of all the non-None values in the list.794 replace_nones_with(val)795 Replaces all the None values in the list with the given value.796 __len__()797 Returns the length of the list.798 __mul__(scalar)799 Multiplies all the not-None values in the list by the given scalar.800 """801 def __init__(self, iterable=()):802 """803 Initializes a ListTracker instance with an empty list to store values.804 """805 self.collection = list(iterable)806 807 def append(self, val):808 """809 Appends a value to the list.810 811 Parameters812 ----------813 val814 The value to append to the list.815 """816 self.collection.append(val)817 818 def sum(self):819 """820 Returns the sum of all the not-None values in the list.821 """ 822 s = 0823 for val in self.collection:824 if val is not None:825 s += val826 return s827 828 def min(self):829 """830 Returns the minimum of all the not-None values in the list.831 832 If the list contains only None values, returns 0.833 """834 min_val = float('inf')835 for val in self.collection:836 if val is not None:837 min_val = min(val, min_val)838 if min_val == float('inf'):839 return 0840 return min_val841 842 def max(self):843 """844 Returns the maximum of all the not-None values in the list.845 846 If the list contains only None values, returns 1.847 """848 max_val = float('-inf')849 for val in self.collection:850 if val is not None:851 max_val = max(val, max_val)852 if max_val == float('-inf'):853 return 1854 return max_val855 856 def count_none(self):857 """858 Returns the number of None values in the list.859 860 This method counts and returns the number of elements in the list861 that are equal to None.862 """863 return self.collection.count(None)864 865 def count_not_none(self):866 """867 Returns the number of non-None values in the list.868 869 This method counts and returns the number of elements in the list870 that are not equal to None.871 """872 cnt = 0873 for val in self.collection:874 if val is not None:875 cnt += 1876 return cnt877 878 def mean(self):879 """880 Returns the mean of all the non-None values in the list.881 882 If the list contains only None values, returns 0.5.883 """884 if self.count_not_none == 0:885 return None886 return self.sum() / self.count_not_none()887 888 def range(self):889 """890 Returns the range of all the non-None values in the list.891 892 If the list contains only None values, returns 1.893 894 Returns895 -------896 float897 The range of the non-None values in the list.898 """899 if self.count_not_none == 0:900 return 1901 return self.max() - self.min()902 903 def replace_nones_with(self, val):904 """905 Replaces all None values in the list with a given value.906 907 This method iterates the list and replaces all the None values with908 the given value.909 910 Parameters911 ----------912 val : Any913 The value to replace None with.914 """915 self.collection = [val if x is None else x for x in self.collection]916 917 def __len__(self):918 """919 Returns the length of the list.920 """921 return len(self.collection)922 923 def __add__(self, scalar):924 return ListTracker([925 val + scalar if val is not None926 else None927 for val in self.collection928 ])929 930 def __radd__(self, scalar):931 return self.__add__(scalar)932 933 def __sub__(self, scalar):934 return ListTracker([935 val - scalar if val is not None936 else None937 for val in self.collection938 ])939 940 def __rsub__(self, scalar):941 return self.__sub__(scalar)942 943 def __mul__(self, scalar): 944 return ListTracker([945 val * scalar if val is not None946 else None947 for val in self.collection948 ])949 950 def __rmul__(self, scalar):951 return self.__mul__(scalar)952 953 def __truediv__(self, scalar):954 return ListTracker([955 val / scalar if val is not None956 else None957 for val in self.collection958 ])959 960 def __rtruediv__(self, scalar):961 return self.__truediv__(scalar)962 963 def __iter__(self):964 """965 Iterate over the list and yield each element.966 967 This method is needed to implement the iterable protocol.968 969 Yields970 ------971 Any972 An element of the list.973 """974 for val in self.collection:975 yield val976 977 def __repr__(self):978 return repr(self.collection)979 980 def __str__(self):981 return str(self.collection)982 983 984NEO4J_DEFAULT_RELATIONS = {985 'to-user': {986 'connected': {987 'self_weight': 1,988 'children_processor': {},989 'children_weights': {}990 },991 'feed': {992 'self_weight': 1,993 'children_processor': {},994 'children_weights': {}995 }996 },997 'to-venue': {998 'visited': {999 'self_weight': 1,1000 'children_processor': {1001 "spending": lambda x : x,1002 "videoViewTime": lambda x : x,1003 "experienceRatings": lambda x : x,1004 "timeSpent": lambda x : x,1005 "visitIds": lambda x : len(x),1006 "profileViewTime": lambda x : x1007 },1008 'children_weights': {1009 "spending": 1,1010 "videoViewTime": 1,1011 "experienceRatings": 1,1012 "timeSpent": 1,1013 "visitIds": 1,1014 "profileViewTime": 11015 }1016 }1017 }1018}1019 1020class Neo4jAggregator(GraphAggregator):1021 def __init__(1022 self,1023 url: str,1024 *args,1025 **kwargs1026 ):1027 """1028 Initialize a Neo4jAggregator object.1029 1030 Parameters1031 ----------1032 url : str1033 The URL of the Neo4j database API endpoint.1034 *args : tuple1035 Additional positional arguments to pass to the GraphAggregator constructor.1036 **kwargs : dict1037 Additional keyword arguments to pass to the GraphAggregator constructor.1038 """1039 super().__init__(*args, **kwargs)1040 self.url = url1041 1042 def get_json(self, path):1043 path = path.strip('/')1044 response = requests.get(self.url.strip('/') + '/' + path)1045 if response.status_code != 200:1046 raise Exception(f"Failed to fetch data from Neo4j API: {response.text}")1047 return response.json()1048 1049 def format_user_data(self, data):1050 text = (f'{data["firstName"]} {data["lastName"]}\n' +1051 f'Bio: {data["bio"]}\n' +1052 f'Occupation: {data["occupation"]}\n' +1053 f'Gender: {data["gender"]}'1054 )1055 return text1056 1057 def format_venue_data(self, data):1058 text = (f'{data["name"]}\n' +1059 f'Description: {data["description"]}\n' +1060 f'Categories: {", ".join(data["category"])}\n' +1061 f'Price range: {data["priceRange"]}'1062 )1063 return text1064 1065 def collect(1066 self,1067 sentence_model: str = 'sentence-transformers/all-mpnet-base-v2',1068 relationships = NEO4J_DEFAULT_RELATIONS1069 ):1070 text_model = SentenceTransformer(sentence_model)1071 user_ids = []1072 1073 for data in tqdm(self.get_json('/entities/user/getall'), desc='Collecting users'):1074 user_ids.append(data['userId'])1075 emb = text_model.encode(self.format_user_data(data))1076 self.add_vertex(1077 vertex_id=data['userId'],1078 category_id='user',1079 data=emb1080 )1081 1082 for data in tqdm(self.get_json('/entities/location/getall'), desc='Collecting venues'):1083 emb = text_model.encode(self.format_venue_data(data))1084 self.add_vertex(1085 vertex_id=data['eventId'],1086 category_id=data['category'][0],1087 data=emb1088 )1089 1090 for uid in tqdm(user_ids, desc='Collecting user-venue edge relationships'):1091 venue_scores = self.calculate_scores(1092 user_id=uid,1093 trg_id_extractor=lambda x: x['location']['eventId'],1094 relationships=relationships['to-venue']1095 )1096 for venue_id, score in venue_scores.items():1097 self.add_edge(1098 src_id=uid,1099 trg_id=venue_id,1100 data=score1101 )1102 1103 for uid in tqdm(user_ids, desc='Collecting user-user edge relationships'):1104 user_scores = self.calculate_scores(1105 user_id=uid,1106 trg_id_extractor=lambda x: x['user']['userId'],1107 relationships=relationships['to-user']1108 )1109 for friend_id, score in user_scores.items():1110 self.add_edge(1111 src_id=uid,1112 trg_id=friend_id,1113 data=score1114 )1115 1116 def calculate_scores(self, user_id, trg_id_extractor, relationships):1117 scores = {} # {trg_id: score}1118 bias = 0 # moving mean of previous scores1119 n_updates = 0 # number of updates1120 relational_normaliser = sum( # normaliser to divide relational weights by1121 [r['self_weight']1122 for r in relationships.values()]1123 )1124 1125 for rel in relationships:1126 1127 rel_config = relationships[rel]1128 rel_weight = rel_config['self_weight'] / relational_normaliser1129 child_normaliser = sum(rel_config['children_weights'].values())1130 1131 endpoint = f'/relationships/{rel}/get-all-outbound-connections/{user_id}'1132 neo4j_data = self.get_json(endpoint)1133 1134 for child, proc, w in zip(1135 rel_config['children_processor'].keys(),1136 rel_config['children_processor'].values(),1137 rel_config['children_weights'].values()1138 ):1139 child_weight = w / child_normaliser1140 total_weight = rel_weight * child_weight1141 trg_ids = []1142 values = ListTracker()1143 1144 for target in neo4j_data:1145 trg_id = trg_id_extractor(target)1146 trg_ids.append(trg_id)1147 1148 if child not in target['relationship']:1149 val = None1150 else:1151 val = proc(target['relationship'][child])1152 values.append(val)1153 1154 values.replace_nones_with(values.mean())1155 if values.range() == 0:1156 if values.max() != 0:1157 values = values - values.max()1158 else:1159 values = (values - values.min()) / (values.range())1160 values = values * total_weight1161 1162 for trg_id, val in zip(trg_ids, values):1163 if trg_id not in scores:1164 scores[trg_id] = bias1165 scores[trg_id] += val1166 1167 bias = (bias * n_updates + values.mean()) / (n_updates + 1)1168 n_updates += 11169 1170 if len(scores) == 0:1171 return scores1172 1173 max_val = max(scores.values())1174 scores = {k: v / max_val for k, v in scores.items()}1175 return scores1176 1177class TestFirebaseAggregator(GraphAggregator):1178 def __init__(1179 self,1180 db,1181 *args,1182 **kwargs1183 ):1184 """1185 Initialize a FirebaseAggregator object.1186 1187 Parameters1188 ----------1189 db : Firestore Client1190 The Firebase database client.1191 *args : tuple1192 Additional positional arguments to pass to the GraphAggregator constructor.1193 **kwargs : dict1194 Additional keyword arguments to pass to the GraphAggregator constructor.1195 """1196 super().__init__(*args, **kwargs)1197 self.db = db1198 1199 def collect(1200 self,