glam4cm 0.1.0__py3-none-any.whl

glam4cm/graph2str/ontouml.py

@@ -0,0 +1,141 @@
+import networkx as nx
+import random
+import itertools
+from tqdm.auto import tqdm
+from constants import *
+from common import (
+    get_node_neighbours,
+    remove_extra_spaces,
+    has_neighbours_incl_incoming
+)
+
+
+def get_node_text_triples(g, distance=1, only_name=False):
+    node_strings = [get_node_str(g, node, distance) for node in g.nodes]
+    node_triples = list()
+    for node, node_str in zip(list(g.nodes), node_strings):
+        name = g.nodes[node]['name'] if 'name' in g.nodes[node] else " reference "
+        node_type = g.nodes[node]['type']
+        prompt_str = f"{node_type} {name}: {node_str}" if not only_name else f"{name}"
+        node_triples.append(prompt_str)
+    return node_triples
+
+
+def check_stereotype_relevance(g, n):
+    return 'use_stereotype' in g.nodes[n] and g.nodes[n]['use_stereotype']
+
+
+def process_name_and_steroetype(g, n):
+    string = g.nodes[n]['name'] if g.nodes[n]['name'] != "Null" else ""
+    string += f' {g.nodes[n]["stereotype"]} ' if check_stereotype_relevance(g, n) else ""
+        
+    return string
+
+
+def process_node_for_string(g, n, src=True):
+    if g.nodes[n]['type'] == 'Class':
+        return [process_name_and_steroetype(g, n)]
+        
+    strings = list()
+    node_str = process_name_and_steroetype(g, n)
+    edges = list(g.in_edges(n)) if src else list(g.out_edges(n))
+    for edge in edges:
+        v = edge[0] if src else edge[1]
+        v_str = f" {process_edge_for_string(g, edge)} {process_name_and_steroetype(g, v)}"
+        n_str = v_str + node_str if src else node_str + v_str
+        strings.append(n_str)
+    return list(set(map(remove_extra_spaces, strings)))
+
+
+def process_edge_for_string(g, e):
+    edge_type_s = e_s[g.edges()[e]['type']]
+    return remove_extra_spaces(f" {edge_type_s} ")
+
+
+def get_triples_from_edges(g, edges=None):
+    if edges is None:
+        edges = g.edges()
+    triples = []
+    for edge in edges:
+        u, v = edge
+        edge_str = process_edge_for_string(g, edge)
+        u_strings, v_strings = process_node_for_string(g, u, src=True), process_node_for_string(g, v, src=False)
+        for u_str, v_str in itertools.product(u_strings, v_strings):
+            pos_triple = u_str + f" {edge_str} " + v_str
+            triples.append(remove_extra_spaces(pos_triple))
+
+    return triples
+
+
+def process_path_string(g, path):
+    edges = list(zip(path[:-1], path[1:]))
+    triples = get_triples_from_edges(g, edges)
+    
+    return remove_extra_spaces(f" {SEP} ".join(triples))
+
+
+def get_triples_from_node(g, n, distance=1):
+    triples = list()
+    use_stereotype = g.nodes[n]['use_stereotype'] if 'use_stereotype' in g.nodes[n] else False
+    g.nodes[n]['use_stereotype'] = False
+    node_neighbours = get_node_neighbours(g, n, distance)
+    for neighbour in node_neighbours:
+        paths = [p for p in nx.all_simple_paths(g, n, neighbour, cutoff=distance)]
+        for path in paths:
+            triples.append(process_path_string(g, path))
+    
+    g.nodes[n]['use_stereotype'] = use_stereotype
+    return triples
+
+
+def get_node_str(g, n, distance=1):
+    node_triples = get_triples_from_node(g, n, distance)
+    return remove_extra_spaces(f" | ".join(node_triples))
+
+
+def create_triples_from_graph_edges(graphs):
+    triples = list()
+    for g, _ in tqdm(graphs):
+        triples += get_triples_from_edges(g)
+
+    return triples
+
+
+def mask_graph(graph, stereotypes_classes, mask_prob=0.2, use_stereotypes=False, use_rel_stereotypes=False):
+    all_stereotype_nodes = [node for node in graph.nodes if 'stereotype' in graph.nodes[node]\
+         and graph.nodes[node]['stereotype'] in stereotypes_classes and has_neighbours_incl_incoming(graph, node)\
+            and (True if use_rel_stereotypes else graph.nodes[node]['type'] == 'Class')]
+    
+    assert all(['stereotype' in graph.nodes[node] for node in all_stereotype_nodes]), "All stereotype nodes should have stereotype property"
+
+    total_masked_nodes = int(len(all_stereotype_nodes) * mask_prob)
+    masked_nodes = random.sample(all_stereotype_nodes, total_masked_nodes)
+    unmasked_nodes = [node for node in all_stereotype_nodes if node not in masked_nodes]
+
+    for node in masked_nodes:
+        graph.nodes[node]['masked'] = True
+        graph.nodes[node]['use_stereotype'] = False
+    
+    for node in unmasked_nodes:
+        graph.nodes[node]['masked'] = False
+        graph.nodes[node]['use_stereotype'] = use_stereotypes
+
+    assert all(['masked' in graph.nodes[node] for node in all_stereotype_nodes]), "All stereotype nodes should be masked or unmasked"
+    
+    
+
+def mask_graphs(graphs, stereotypes_classes, mask_prob=0.2, use_stereotypes=False, use_rel_stereotypes=False):
+    masked, unmasked, total = 0, 0, 0
+    # for graph, f_name in tqdm(graphs, desc='Masking graphs'):
+    for graph, _ in graphs:
+        mask_graph(graph, stereotypes_classes, mask_prob=mask_prob, use_stereotypes=use_stereotypes, use_rel_stereotypes=use_rel_stereotypes)
+        masked += len([node for node in graph.nodes if 'masked' in graph.nodes[node] and graph.nodes[node]['masked']])
+        unmasked += len([node for node in graph.nodes if 'masked' in graph.nodes[node] and not graph.nodes[node]['masked']])
+        total += len([node for node in graph.nodes if 'masked' in graph.nodes[node]])
+        
+    ## % of masked nodes upto 2 decimal places
+    print(f"Masked {round(masked/total, 2)*100}%")
+    print(f"Unmasked {round(unmasked/total, 2)*100}%")
+
+    print("Total masked nodes:", masked)
+    print("Total unmasked nodes:", unmasked)

glam4cm/lang2graph/archimate.py

@@ -0,0 +1,31 @@
+from glam4cm.lang2graph.common import LangGraph
+
+
+class ArchiMateNxG(LangGraph):
+    def __init__(
+        self, 
+        json_obj: dict, 
+        path: str,
+        timeout = -1
+    ):
+        super().__init__()
+        self.json_obj = json_obj
+        self.timeout = timeout
+        self.path = path
+        self.graph_id = json_obj['identifier'].split('/')[-1]
+
+
+        self.__create_graph()
+        self.set_numbered_labels()
+
+        # self.text = " ".join([
+        #     self.nodes[node]['name'] if 'name' in self.nodes[node] else ''
+        #     for node in self.nodes
+        # ])
+
+
+    def __create_graph(self):
+        for node in self.json_obj['elements']:
+            self.add_node(node['id'], **node)
+        for edge in self.json_obj['relationships']:
+            self.add_edge(edge['sourceId'], edge['targetId'], **edge)

glam4cm/lang2graph/common.py

@@ -0,0 +1,416 @@
+from abc import abstractmethod
+from typing import List
+import networkx as nx
+from uuid import uuid4
+import numpy as np
+import torch
+from glam4cm.data_loading.metadata import GraphMetadata
+from glam4cm.tokenization.special_tokens import *
+from glam4cm.tokenization.utils import doc_tokenizer
+import glam4cm.utils as utils
+
+SEP = ' '
+REFERENCE = 'reference'
+SUPERTYPE = 'supertype'
+CONTAINMENT = 'containment'
+
+
+class LangGraph(nx.DiGraph):
+    def __init__(self):
+        super().__init__()
+        self.id = uuid4().hex
+        self.node_label_to_id = {}
+        self.id_to_node_label = {}
+        self.edge_label_to_id = {}
+        self.id_to_edge_label = {}
+
+
+    @abstractmethod
+    def create_graph(self):
+        pass
+
+
+    def set_numbered_labels(self):
+        self.node_label_to_id = {label: i for i, label in enumerate(self.nodes())}
+        self.id_to_node_label = {i: label for i, label in enumerate(self.nodes())}
+
+        self.edge_label_to_id = {label: i for i, label in enumerate(self.edges())}
+        self.id_to_edge_label = {i: label for i, label in enumerate(self.edges())}
+
+        self.numbered_graph = self.get_numbered_graph()
+        self.edge_to_idx = {edge: idx for idx, edge in enumerate(self.numbered_graph.edges())}
+        self.idx_to_edge = {idx: edge for idx, edge in enumerate(self.numbered_graph.edges())}
+
+
+
+    def get_numbered_graph(self) -> nx.DiGraph:
+        nodes = [(self.node_label_to_id[i], data) for i, data in list(self.nodes(data=True))]
+        edges = [(self.node_label_to_id[i], self.node_label_to_id[j], data) for i, j, data in list(self.edges(data=True))]
+        graph = nx.DiGraph()
+        graph.add_nodes_from(nodes)
+        graph.add_edges_from(edges)
+
+        return graph
+
+
+    @property
+    def enr(self):
+        if self.number_of_nodes() == 0:
+            return -1
+        return self.number_of_edges() / self.number_of_nodes()
+
+
+    @property
+    def edge_index(self):
+        edge_index = torch.tensor(list(self.numbered_graph.edges)).t().contiguous().numpy()
+        return edge_index
+
+    @property
+    def hash(self):
+        return utils.md5_hash(str(sorted(self.edges)))
+    
+    def get_edge_id(self, edge):
+        return self.edge_label_to_id[edge]
+
+    def get_edge_label(self, edge_id):
+        return self.edge_label_to_id[edge_id]
+
+    
+    def get_node_id(self, node):
+        return self.node_label_to_id[node]
+    
+    def get_node_label(self, node_id):
+        return self.node_label_to_id[node_id]
+
+
+def create_graph_from_edge_index(graph, edge_index: np.ndarray):
+    """
+    Create a subgraph from G using only the edges specified in edge_index.
+    
+    Parameters:
+    G (networkx.Graph): The original graph.
+    edge_index (numpy.ndarray): A numpy containing edge indices.
+    
+    Returns:
+    networkx.Graph: A subgraph of G containing only the edges in edge_index.
+    """
+
+    if isinstance(edge_index, torch.Tensor):
+        edge_index = edge_index.cpu().numpy()
+
+    # Add nodes and edges from the edge_index to the subgraph
+    subgraph = nx.DiGraph()
+    subgraph.add_nodes_from(list(graph.numbered_graph.nodes(data=True)))
+    subgraph.add_edges_from([(u, v, graph.numbered_graph.edges[u, v]) for u, v in edge_index.T])
+    for node, data in subgraph.nodes(data=True):
+        data = graph.numbered_graph.nodes[node]
+        subgraph.nodes[node].update(data)
+
+
+
+    subgraph.node_label_to_id = graph.node_label_to_id
+    subgraph.id_to_node_label = graph.id_to_node_label
+    subgraph.edge_label_to_id = graph.edge_label_to_id
+    subgraph.id_to_edge_label = graph.id_to_edge_label
+    try:
+        assert subgraph.number_of_edges() == edge_index.shape[1]
+    except AssertionError as e:
+        print(f"Number of edges mismatch {subgraph.number_of_edges()} != {edge_index.size(1)}")
+        import pickle
+        pickle.dump([graph, edge_index], open("subgraph.pkl", "wb"))
+        raise e
+
+    return subgraph
+
+
+def format_path(
+    graph: LangGraph, 
+    path: List, 
+    metadata: GraphMetadata, 
+    use_node_attributes = False, 
+    use_node_types = False, 
+    use_edge_label = False, 
+    use_edge_types = False, 
+    node_cls_label=None,
+    edge_cls_label='type',
+    use_special_tokens = False, 
+    no_labels = False,
+    preprocessor=doc_tokenizer,
+    neg_sample=False
+):
+    """Format a path into a string representation."""
+    def get_node_label(node):
+        masked = graph.nodes[node].get('masked')
+        node_type = f"{graph.nodes[node].get(f'{node_cls_label}', '')}" if use_node_types and not masked and node_cls_label else ''
+        node_type = f"{node_cls_label}: {node_type}" if node_type else ''
+        node_label = get_node_name(
+            graph.nodes[node], 
+            metadata.node_label, 
+            use_node_attributes, 
+            metadata.node_attributes
+        ) if not no_labels else ''
+        if preprocessor:
+            node_label = preprocessor(node_label)
+        
+        node_label = f"{node_type} {node_label}".strip()
+        if use_special_tokens:
+            node_label = f"{NODE_BEGIN}{node_label}{NODE_END}"
+        
+        return node_label.strip()
+
+    def get_edge_label(n1, n2):
+        edge_data = graph.get_edge_data(n1, n2)
+        masked = edge_data.get('masked')
+        edge_label = edge_data.get(metadata.edge_label, '') if use_edge_label and not no_labels else ''
+        edge_type = f"{edge_cls_label}:{get_edge_data(edge_data, f'{edge_cls_label}', metadata.type)}" if use_edge_types and not masked  and edge_cls_label else ''
+        
+        if preprocessor:
+            edge_label = preprocessor(edge_label)
+        
+        edge_label = f"{edge_type} {edge_label}".strip()
+        
+        if use_special_tokens:
+            edge_label = f"{EDGE_START}{edge_label}{EDGE_END}"
+        
+        return edge_label.strip()
+
+    assert len(path) > 0, "Path must contain at least one node."
+    formatted = [get_node_label(path[0])]
+    for i in range(1, len(path)):
+        n1 = path[i - 1]
+        n2 = path[i]
+
+        if not neg_sample:
+            formatted.append(get_edge_label(n1, n2))
+        formatted.append(get_node_label(n2))
+
+    return " ".join(formatted).strip()
+
+def get_edge_texts(
+    graph: LangGraph, 
+    edge: tuple, 
+    d: int, 
+    metadata: GraphMetadata, 
+    use_node_attributes=False, 
+    use_node_types=False, 
+    use_edge_types=False,
+    use_edge_label=False,
+    node_cls_label=None,
+    edge_cls_label='type',
+    use_special_tokens=False, 
+    no_labels=False,
+    preprocessor: callable = doc_tokenizer,
+    neg_samples=False
+):
+    n1, n2 = edge
+    if not neg_samples:
+        masked = graph.edges[n1, n2].get('masked')
+        graph.edges[n1, n2]['masked'] = True
+
+    n1_text = get_node_text(
+        graph=graph,
+        node=n1,
+        d=d,
+        metadata=metadata,
+        use_node_attributes=use_node_attributes,
+        use_node_types=use_node_types,
+        use_edge_types=use_edge_types,
+        use_edge_label=use_edge_label,
+        node_cls_label=node_cls_label,
+        edge_cls_label=edge_cls_label,
+        use_special_tokens=use_special_tokens,
+        no_labels=no_labels,
+        preprocessor=preprocessor,
+        exclude_edges=[edge]
+    )
+    n2_text = get_node_text(
+        graph=graph,
+        node=n2,
+        d=d,
+        metadata=metadata,
+        use_node_attributes=use_node_attributes,
+        use_node_types=use_node_types,
+        use_edge_types=use_edge_types,
+        use_edge_label=use_edge_label,
+        node_cls_label=node_cls_label,
+        edge_cls_label=edge_cls_label,
+        use_special_tokens=use_special_tokens,
+        no_labels=no_labels,
+        preprocessor=preprocessor,
+        exclude_edges=[edge]
+    )
+    if not neg_samples:
+        graph.edges[n1, n2]['masked'] = masked or False
+
+    return n1_text + EDGE_START + EDGE_END + n2_text
+
+
+def get_node_text(
+    graph: LangGraph, 
+    node, 
+    d: int, 
+    metadata: GraphMetadata, 
+    use_node_attributes=False, 
+    use_node_types=False, 
+    use_edge_types=False,
+    use_edge_label=False,
+    node_cls_label=None,
+    edge_cls_label='type',
+    use_special_tokens=False, 
+    no_labels=False,
+    preprocessor: callable = doc_tokenizer,
+    exclude_edges=None
+):
+    masked = graph.nodes[node].get('masked')
+    graph.nodes[node]['masked'] = True
+    raw_paths = utils.bfs(graph=graph, start_node=node, d=d, exclude_edges=exclude_edges)
+    unique_paths = utils.remove_subsets(list_of_lists=raw_paths)
+    text = "\n".join([
+        format_path(
+            graph=graph, 
+            path=path, 
+            metadata=metadata, 
+            use_node_attributes=use_node_attributes, 
+            use_node_types=use_node_types, 
+            use_edge_types=use_edge_types, 
+            use_edge_label=use_edge_label,
+            node_cls_label=node_cls_label,
+            edge_cls_label=edge_cls_label,
+            use_special_tokens=use_special_tokens,
+            no_labels=no_labels,
+            preprocessor=preprocessor, 
+            neg_sample=False
+        )
+        for path in unique_paths
+    ])
+    graph.nodes[node]['masked'] = masked or False
+    return text
+
+
+def get_node_texts(
+        graph: LangGraph, 
+        d: int, 
+        metadata: GraphMetadata, 
+        use_node_attributes=False, 
+        use_node_types=False, 
+        use_edge_types=False,
+        use_edge_label=False,
+        node_cls_label=None,
+        edge_cls_label='type',
+        use_special_tokens=False, 
+        no_labels=False,
+        preprocessor: callable = doc_tokenizer
+    ):
+    paths_dict = {}
+    for node in graph.nodes:
+        paths_dict[node] = get_node_text(
+            graph=graph,
+            node=node,
+            d=d,
+            metadata=metadata,
+            use_node_attributes=use_node_attributes,
+            use_node_types=use_node_types,
+            use_edge_types=use_edge_types,
+            use_edge_label=use_edge_label,
+            node_cls_label=node_cls_label,
+            edge_cls_label=edge_cls_label,
+            use_special_tokens=use_special_tokens,
+            no_labels=no_labels,
+            preprocessor=preprocessor
+        )
+
+    return paths_dict
+
+
+def get_attribute_labels(node_data, attribute_labels):
+    if isinstance(node_data[attribute_labels], list):
+        if not node_data[attribute_labels]:
+            return ''
+        if isinstance(node_data[attribute_labels][0], tuple):
+            return ", ".join([f"{k}: {v}" for k, v in node_data[attribute_labels]])
+        elif isinstance(node_data[attribute_labels][0], dict):
+            return ", ".join([f"{k}: {v}" for d in node_data[attribute_labels] for k, v in d.items()])
+        return ", ".join(node_data[attribute_labels])
+    if isinstance(node_data[attribute_labels], dict):
+        return ", ".join([f"{k}: {v}" for k, v in node_data[attribute_labels].items()])
+    return node_data[attribute_labels]
+
+
+def get_node_name(
+        node_data, 
+        label, 
+        use_attributes,
+        attribute_labels,
+    ):
+    if use_attributes and attribute_labels in node_data:
+        attributes_str = "(" + get_attribute_labels(node_data, attribute_labels) + ")"
+    else:
+        attributes_str = ''
+    node_label = node_data.get(label, '')
+    node_label = '' if node_label.lower() == 'null' else node_label
+    return f"{node_label}{attributes_str}".strip()
+
+
+def get_node_data(
+    node_data: dict,
+    node_label: str,
+    model_type: str,
+):
+    if model_type == 'archimate':
+        return get_archimate_node_data(node_data, node_label)
+    elif model_type == 'ecore':
+        return get_uml_node_data(node_data, node_label)
+    elif model_type == 'ontouml':
+        return get_ontouml_node_data(node_data, node_label)
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+
+
+
+def get_edge_data(
+    edge_data: dict,
+    edge_label: str,
+    model_type: str,
+):
+    if model_type == 'archimate':
+        return get_archimate_edge_data(edge_data, edge_label)
+    elif model_type == 'ecore':
+        return get_uml_edge_data(edge_data, edge_label)
+    elif model_type == 'ontouml':
+        return get_ontouml_edge_data(edge_data, edge_label)
+    else:
+        raise ValueError(f"Unknown model type: {model_type}")
+
+
+def get_archimate_node_data(edge_data: dict, node_label: str):
+    return edge_data.get(node_label)
+
+def get_uml_node_data(node_data: dict, node_label: str):
+    return node_data.get(node_label, '')
+
+def get_ontouml_node_data(node_data: dict, node_label: str):
+    return node_data.get(node_label, '')
+
+
+def get_archimate_edge_data(edge_data: dict, edge_label: str):
+    return edge_data.get(edge_label)
+
+
+def get_uml_edge_data(edge_data: dict, edge_label: str):
+    if edge_label == 'type':
+        return get_uml_edge_type(edge_data)
+    elif edge_label in edge_data:
+        return edge_data[edge_label]
+    else:
+        raise ValueError(f"Unknown edge label: {edge_label}")
+
+def get_ontouml_edge_data(edge_data: dict, edge_label: str):
+    return edge_data.get(edge_label)
+
+def get_uml_edge_type(edge_data):
+    edge_type = edge_data.get('type')
+    if edge_type == SUPERTYPE:
+        return SUPERTYPE
+    if edge_type == CONTAINMENT:
+        return CONTAINMENT
+    return REFERENCE