glam4cm 0.1.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

glam4cm/downstream_tasks/gnn_link_prediction.py

@@ -1,14 +1,18 @@
 import os
 from glam4cm.data_loading.graph_dataset import GraphEdgeDataset
 from glam4cm.models.gnn_layers import GNNConv, EdgeClassifer
-from glam4cm.settings import LP_TASK_LINK_PRED
+from glam4cm.settings import LINK_PRED_TASK, results_dir
 from glam4cm.downstream_tasks.utils import get_models_dataset
 from glam4cm.tokenization.special_tokens import *
 from glam4cm.trainers.gnn_link_predictor import GNNLinkPredictionTrainer as Trainer
-from glam4cm.utils import merge_argument_parsers, set_seed
-from glam4cm.downstream_tasks.common_args import get_common_args_parser, get_config_params, get_gnn_args_parser
-
+from glam4cm.utils import merge_argument_parsers, set_seed, set_torch_encoding_labels
+from glam4cm.downstream_tasks.common_args import (
+    get_common_args_parser, 
+    get_config_params, 
+    get_gnn_args_parser
+)
 
+ 
 def get_parser():
     common_parser = get_common_args_parser()
     gnn_parser = get_gnn_args_parser()
@@ -21,6 +25,7 @@ def run(args):
     set_seed(args.seed)
     
     config_params = dict(
+        include_dummies = args.include_dummies,
         min_enr = args.min_enr,
         min_edges = args.min_edges,
         remove_duplicates = args.remove_duplicates,
@@ -42,14 +47,18 @@ def run(args):
     aggregation = args.aggregation
 
     graph_data_params = get_config_params(args)
+    
+    if args.use_embeddings:
+        graph_data_params['embed_model_name'] = os.path.join(results_dir, dataset_name, f"LM_{LINK_PRED_TASK}")
+    
     print("Loading graph dataset")
     graph_dataset = GraphEdgeDataset(
-        dataset, 
-        dict(
+        dataset,
+        task_type=LINK_PRED_TASK, 
+        **dict(
             **graph_data_params, 
-            add_negative_train_samples=args.add_negative_train_samples, 
+            add_negative_train_samples=True, 
             neg_sampling_ratio=args.neg_sampling_ratio,
-            task=LP_TASK_LINK_PRED
     ))
 
     input_dim = graph_dataset[0].data.x.shape[1]
@@ -78,7 +87,7 @@ def run(args):
     logs_dir = os.path.join(
         "logs",
         dataset_name,
-        "gnn_lp",
+        f"GNN_{LINK_PRED_TASK}",
         f'{graph_dataset.config_hash}',
     )
 
@@ -92,11 +101,14 @@ def run(args):
         bias=False,
     )
 
+    graph_torch_data = graph_dataset.get_torch_dataset()
+    # exclude_labels = getattr(graph_dataset, f"node_exclude_{args.node_cls_label}")
+    # set_torch_encoding_labels(graph_torch_data, f"node_{args.node_cls_label}", exclude_labels)
     
     trainer = Trainer(
-        gnn_conv_model, 
-        mlp_predictor, 
-        graph_dataset.get_torch_dataset(),
+        model=gnn_conv_model, 
+        predictor=mlp_predictor, 
+        dataset=graph_torch_data,
         lr=args.lr,
         num_epochs=args.num_epochs,
         batch_size=args.batch_size,
@@ -106,4 +118,4 @@ def run(args):
 
 
     print("Training GNN Link Prediction model")
-    trainer.run()
+    trainer.run()

glam4cm/downstream_tasks/gnn_node_classification.py

@@ -2,10 +2,15 @@ import os
 from glam4cm.data_loading.graph_dataset import GraphNodeDataset
 from glam4cm.models.gnn_layers import GNNConv, NodeClassifier
 from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.settings import NODE_CLS_TASK, results_dir
 from glam4cm.tokenization.special_tokens import *
 from glam4cm.trainers.gnn_node_classifier import GNNNodeClassificationTrainer as Trainer
-from glam4cm.utils import merge_argument_parsers, set_seed
-from glam4cm.downstream_tasks.common_args import get_common_args_parser, get_config_params, get_gnn_args_parser
+from glam4cm.utils import merge_argument_parsers, set_seed, set_torch_encoding_labels
+from glam4cm.downstream_tasks.common_args import (
+    get_common_args_parser, 
+    get_config_params, 
+    get_gnn_args_parser
+)
 
 
 def get_parser():
@@ -20,6 +25,7 @@ def run(args):
     set_seed(args.seed)
     
     config_params = dict(
+        include_dummies = args.include_dummies,
         min_enr = args.min_enr,
         min_edges = args.min_edges,
         remove_duplicates = args.remove_duplicates,
@@ -29,13 +35,19 @@ def run(args):
     dataset_name = args.dataset
 
     dataset = get_models_dataset(dataset_name, **config_params)
-    graph_data_params = get_config_params(args)
+    graph_data_params = {**get_config_params(args), 'task_type': NODE_CLS_TASK}
+    
+    if args.use_embeddings:
+        graph_data_params['embed_model_name'] = os.path.join(results_dir, dataset_name, f'{args.node_cls_label}')
 
     print("Loading graph dataset")
     graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
     print("Loaded graph dataset")
-
+    
+    
     graph_torch_data = graph_dataset.get_torch_dataset()
+    exclude_labels = getattr(graph_dataset, f"node_exclude_{args.node_cls_label}")
+    set_torch_encoding_labels(graph_torch_data, f"node_{args.node_cls_label}", exclude_labels)
 
     num_nodes_label = f"num_nodes_{args.node_cls_label}"
     assert hasattr(graph_dataset, num_nodes_label), f"Graph dataset does not have attribute {num_nodes_label}"
@@ -83,7 +95,7 @@ def run(args):
     logs_dir = os.path.join(
         "logs",
         dataset_name,
-        "gnn_node_cls",
+        f"GNN_{NODE_CLS_TASK}",
         f"{graph_dataset.config_hash}",
     )
 
@@ -92,7 +104,7 @@ def run(args):
         mlp_predictor, 
         graph_torch_data,
         cls_label=args.node_cls_label,
-        exclude_labels=getattr(graph_dataset, f"node_exclude_{args.node_cls_label}"),
+        exclude_labels=[-1],
         lr=args.lr,
         num_epochs=args.num_epochs,
         use_edge_attrs=args.use_edge_attrs,
@@ -100,4 +112,4 @@ def run(args):
     )
 
     print("Training GNN Node Classification model")
-    trainer.run()
+    trainer.run()

glam4cm/downstream_tasks/utils.py

@@ -1,6 +1,7 @@
 from glam4cm.data_loading.models_dataset import (
     ArchiMateDataset, 
-    EcoreDataset
+    EcoreDataset,
+    OntoUMLDataset
 )
 
 
@@ -8,7 +9,8 @@ dataset_to_metamodel = {
     'modelset': 'ecore',
     'ecore_555': 'ecore',
     'mar-ecore-github': 'ecore',
-    'eamodelset': 'ea'
+    'eamodelset': 'ea',
+    'ontouml': 'ontouml',
 }
 
 
@@ -22,6 +24,8 @@ def get_model_dataset_class(dataset_name):
         dataset_class = ArchiMateDataset
     elif dataset_type == 'ecore':
         dataset_class = EcoreDataset
+    elif dataset_type == 'ontouml':
+        dataset_class = OntoUMLDataset
     else:
         raise ValueError(f"Unknown dataset type: {dataset_type}")
     return dataset_class
@@ -33,3 +37,13 @@ def get_models_dataset(dataset_name, **config_params):
         del config_params['language']
     dataset_class = get_model_dataset_class(dataset_name)
     return dataset_class(dataset_name, **config_params)
+
+
+def get_logging_steps(dataset_size, num_epochs, batch_size):
+    """
+    Calculate the logging steps based on the dataset size, number of epochs, and batch size.
+    """
+    num_steps = dataset_size // batch_size
+    logging_steps = num_steps * num_epochs // 20
+    print(f"Logging steps: {logging_steps}")
+    return logging_steps

glam4cm/embeddings/bert.py

@@ -36,7 +36,7 @@ class BertEmbedder(Embedder):
         print("Number of Texts: ", len(text))
 
         dataset = EncodingDataset(self.tokenizer, texts=text, remove_duplicates=False)
-        loader = DataLoader(dataset, batch_size=256)
+        loader = DataLoader(dataset, batch_size=64)
 
         embeddings = list()
         with torch.no_grad():

glam4cm/embeddings/common.py

@@ -5,7 +5,9 @@ from typing import List, Union
 import torch
 from glam4cm.settings import (
     WORD2VEC_MODEL,
-    TFIDF_MODEL
+    TFIDF_MODEL,
+    MODERN_BERT, 
+    BERT_MODEL
 )
 
 
@@ -27,10 +29,11 @@ def get_embedding_model(
         model_name: str,
         ckpt: str = None
     ) -> Embedder:
-    if ckpt:
-        model_name = json.load(open(os.path.join(ckpt, 'config.json')))['_name_or_path']
+    # if ckpt:
+    #     model_name = json.load(open(os.path.join(ckpt, 'config.json')))['_name_or_path']
+    #     print("Model name:", model_name)
         
-    if 'bert' in model_name:
+    if model_name in [MODERN_BERT, BERT_MODEL]:
         from glam4cm.embeddings.bert import BertEmbedder
         return BertEmbedder(model_name, ckpt)
     elif WORD2VEC_MODEL in model_name:

glam4cm/encoding/encoders.py

@@ -4,7 +4,7 @@ from sklearn.preprocessing import LabelEncoder
 import fasttext
 from scipy.sparse import csr_matrix
 import numpy as np
-from encoding.common import (
+from glam4cm.encoding.common import (
     doc_tokenizer, 
     SEP
 )

glam4cm/lang2graph/archimate.py

@@ -18,11 +18,6 @@ class ArchiMateNxG(LangGraph):
         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']:

glam4cm/lang2graph/common.py

@@ -8,11 +8,15 @@ 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
+from glam4cm.settings import (
+    SUPERTYPE,
+    REFERENCE,
+    CONTAINMENT,
+    
+    EDGE_CLS_TASK,
+    LINK_PRED_TASK,
+)
 
-SEP = ' '
-REFERENCE = 'reference'
-SUPERTYPE = 'supertype'
-CONTAINMENT = 'containment'
 
 
 class LangGraph(nx.DiGraph):
@@ -112,13 +116,14 @@ def create_graph_from_edge_index(graph, edge_index: np.ndarray):
     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
+    if len(edge_index) > 0:
+        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
 
@@ -140,15 +145,24 @@ def format_path(
 ):
     """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_type = f"{graph.nodes[node].get(f'{node_cls_label}', '')}" \
+            if use_node_types and not masked and node_cls_label else ''    
+        
+        if node_type != '':
+            if isinstance(graph.nodes[node].get(f'{node_cls_label}'), bool):
+                node_type = node_cls_label.title() if graph.nodes[node].get(f'{node_cls_label}') 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)
         
@@ -174,8 +188,9 @@ def format_path(
         
         return edge_label.strip()
 
+    # import code; code.interact(local=locals())
     assert len(path) > 0, "Path must contain at least one node."
-    formatted = [get_node_label(path[0])]
+    formatted = []
     for i in range(1, len(path)):
         n1 = path[i - 1]
         n2 = path[i]
@@ -184,12 +199,18 @@ def format_path(
             formatted.append(get_edge_label(n1, n2))
         formatted.append(get_node_label(n2))
 
-    return " ".join(formatted).strip()
+    node_str = get_node_label(path[0])
+    if len(formatted) > 0:
+        node_str += " | " + " ".join(formatted).strip()
+    
+    return node_str
+
 
 def get_edge_texts(
     graph: LangGraph, 
     edge: tuple, 
     d: int, 
+    task_type: str,
     metadata: GraphMetadata, 
     use_node_attributes=False, 
     use_node_types=False, 
@@ -206,7 +227,8 @@ def get_edge_texts(
     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,
@@ -239,10 +261,26 @@ def get_edge_texts(
         preprocessor=preprocessor,
         exclude_edges=[edge]
     )
+    
+
+    edge_text = ""    
+    
     if not neg_samples:
         graph.edges[n1, n2]['masked'] = masked or False
+        
+        edge_data = graph.get_edge_data(n1, n2)
+        edge_type = get_edge_data(edge_data, edge_cls_label, metadata.type)
+        edge_label = edge_data.get(metadata.edge_label, '') if use_edge_label and not no_labels else ''
+        
+        if task_type not in [EDGE_CLS_TASK, LINK_PRED_TASK]:
+            if use_edge_types :
+                edge_text += f" {edge_cls_label}: {edge_type} " if not no_labels else ''
+                
+            if use_edge_label:
+                edge_text += f" {edge_label} " if not no_labels else ''
 
-    return n1_text + EDGE_START + EDGE_END + n2_text
+    
+    return n1_text + EDGE_START + f"{edge_text}" + EDGE_END + n2_text
 
 
 def get_node_text(
@@ -263,28 +301,39 @@ def get_node_text(
 ):
     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
-    ])
+    # 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)
+    node_neighbour_texts = list()
+    node_neighbours = utils.get_node_neighbours(graph, node, d, exclude_edges=exclude_edges)
+    for neighbour in node_neighbours:
+        unique_paths = [p for p in nx.all_simple_paths(graph, node, neighbour, cutoff=d)]
+
+        node_neighbour_texts.extend([
+            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
+    node_str = "\n".join(node_neighbour_texts).strip() if node_neighbour_texts else ''
+    
+    if node_cls_label == 'stereotype':
+        node_str = graph.nodes[node]['type'].title() + " " + node_str
+    
+    return node_str.strip()
 
 
 def get_node_texts(
@@ -326,6 +375,8 @@ 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], str):
+            return ", ".join(node_data[attribute_labels])
         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):
@@ -346,8 +397,12 @@ def get_node_name(
         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
+    
+    node_label = node_data.get(label, '') if node_data.get(label, '') else ''
+    node_label = '' if node_label and node_label.lower() in ['null', 'none'] else node_label
+    # if attributes_str:
+    #     print(f"Node label: {node_label} | Attributes: {attributes_str}")
+        
     return f"{node_label}{attributes_str}".strip()
 
 
@@ -405,7 +460,10 @@ def get_uml_edge_data(edge_data: dict, edge_label: str):
         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)
+    try:
+        return {'rel': "relates", "gen": "generalizes"}[edge_data.get(edge_label)]
+    except KeyError:
+        raise ValueError(f"Unknown edge label: {edge_label}")
 
 def get_uml_edge_type(edge_data):
     edge_type = edge_data.get('type')

glam4cm/lang2graph/ecore.py

@@ -58,8 +58,7 @@ class EcoreNxG(LangGraph):
             for f in structural_features:
                 if f['type'] == 'ecore:EAttribute':
                     name = f['name']
-                    attr_type = f['ref'] if f['ref'] else ''
-                    attributes.append((name, attr_type))
+                    attributes.append(name)
 
             self.add_node(
                 classifier_name, 

glam4cm/lang2graph/ontouml.py

@@ -48,13 +48,14 @@ extra_properties = [
 class OntoUMLNxG(LangGraph):
     def __init__(self, json_obj: dict, rel_as_node=True):
         super().__init__()
+        self.graph_id = json_obj['id']
         self.json_obj = json_obj
         self.rel_as_node = rel_as_node
         self.__create_graph()
         self.set_numbered_labels()
 
         self.text = " ".join([
-            self.nodes[node]['name'] if 'name' in self.nodes[node] else ''
+            self.nodes[node]['name'] if 'name' in self.nodes[node] and self.nodes[node]['name'] else ''
             for node in self.nodes
         ])
     
@@ -76,6 +77,7 @@ class OntoUMLNxG(LangGraph):
                             ontouml_id2obj(item)
 
         def create_nxg():
+
             for k, v in id2obj_map.items():
                 node_name = v.get('name', '')
                 
@@ -85,7 +87,8 @@ class OntoUMLNxG(LangGraph):
                         self.nodes[k][prop] = v[prop] if prop in v else False
 
                     logger.info(f"Node: {node_name} type: {v[ONTOUML_ELEMENT_TYPE]}")
-
+                # else:
+                #     continue
                 
                 logger.info(f"Node: {node_name} type: {v[ONTOUML_ELEMENT_TYPE]}")
                 if ONTOUML_STEREOTYPE in v and v[ONTOUML_STEREOTYPE] is not None:
@@ -108,10 +111,8 @@ class OntoUMLNxG(LangGraph):
                     
                     elif ONTOUML_PROPERTIES in v and v[ONTOUML_PROPERTIES] is not None:
                         properties = v[ONTOUML_PROPERTIES] if isinstance(v[ONTOUML_PROPERTIES], list) else [v[ONTOUML_PROPERTIES]]
-                        properties_str = ", ".join([property[ONTOUML_ELEMENT_NAME] for property in properties])
-                        self.nodes[k][ONTOUML_PROPERTIES] = properties_str
-                        logger.info(f"Properties: {properties_str}")
-
+                        self.nodes[k][ONTOUML_PROPERTIES] = [property[ONTOUML_ELEMENT_NAME] for property in properties]
+                        
 
                 elif v[ONTOUML_ELEMENT_TYPE] == ONTOUML_RELATION:    
                     properties = v[ONTOUML_PROPERTIES] if isinstance(v[ONTOUML_PROPERTIES], list) else [v[ONTOUML_PROPERTIES]]
@@ -144,7 +145,7 @@ class OntoUMLNxG(LangGraph):
 
         def create_nxg_rel_as_edge():
             # TODO: To be implemented
-            pass
+            raise NotImplementedError
 
 
         id2obj_map = dict()

glam4cm/models/gnn_layers.py

@@ -123,19 +123,33 @@ class GNNConv(torch.nn.Module):
                 h = self.dropout(h)
             return h
 
+        edge_attr_val = isinstance(edge_attr, torch.Tensor) and self.is_headed_model()
         h = in_feat
-        h = self.conv_layers[0](h, edge_index, edge_attr) if isinstance(edge_attr, torch.Tensor) else self.conv_layers[0](h, edge_index)
-        activate(h)
+        h = self.conv_layers[0](h, edge_index, edge_attr) \
+        if edge_attr_val else self.conv_layers[0](h, edge_index)
+        h = activate(h)
 
         for conv in self.conv_layers[1:-1]:
-            nh = conv(h, edge_index, edge_attr) if isinstance(edge_attr, torch.Tensor) else conv(h, edge_index)
+            nh = conv(h, edge_index, edge_attr) if edge_attr_val else conv(h, edge_index)
             h = nh if not self.residual else nh + h
-            activate(h)
+            h = activate(h)
         
         h = self.conv_layers[-1](h, edge_index)
-        activate(h)
+        h = activate(h)
         return h
-  
+
+    def is_headed_model(self):
+        """"
+        Returns True if the model is a headed model
+        Checks if the model name is in the supported_conv_models dictionary
+        and if the model requires num_heads
+        """
+        headed = self.num_heads is not None
+        model_name = self.conv_layers[0].__class__.__name__
+        if model_name in supported_conv_models:
+            return supported_conv_models[model_name] and headed
+        return False
+        
 
 class EdgeClassifer(nn.Module):
 

glam4cm/models/hf.py

@@ -1,7 +1,7 @@
 from transformers import AutoModelForSequenceClassification
 
-def get_model(model_name, num_labels, len_tokenizer=None) -> AutoModelForSequenceClassification:
-    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels)
+def get_model(model_name, num_labels, len_tokenizer=None, trust_remote_code=False) -> AutoModelForSequenceClassification:
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=num_labels, trust_remote_code=trust_remote_code)
     if len_tokenizer:
         model.resize_token_embeddings(len_tokenizer)
         assert model.config.vocab_size == len_tokenizer,\

glam4cm/run.py

@@ -43,8 +43,8 @@ tasks = {
     
     6: 'GNN Graph Classification',
     7: 'GNN Node Classification',
-    8: 'GNN Edge Classification',
-    9: 'GNN Link Prediction',
+    8: 'GNN Link Prediction',
+    9: 'GNN Edge Classification',
     10: 'CM-GPT Causal Modeling',
     11: 'CM-GPT Node Classification',
     12: 'CM-GPT Edge Classification'
@@ -60,16 +60,15 @@ tasks_handler_map = {
     5: (bert_edge_classification.run, bert_ec_parse_args),
     6: (gnn_graph_cls.run, gnn_parse_args),
     7: (gnn_node_classification.run, gnn_nc_parse_args),
-    8: (gnn_edge_classification.run, gnn_ec_parse_args),
-    9: (gnn_link_prediction.run, gnn_lp_parse_args),
+    8: (gnn_link_prediction.run, gnn_lp_parse_args),
+    9: (gnn_edge_classification.run, gnn_ec_parse_args),
     10: (cm_gpt_pretraining.run, cm_gpt_parse_args),
     11: (cm_gpt_node_classification.run, cm_gpt_nc_parse_args),
     12: (cm_gpt_edge_classification.run, cm_gpt_ec_parse_args)
 }
 
 
-if __name__ == '__main__':
-
+def main():
     main_parser = argparse.ArgumentParser(description="Train ML models on conceptual models")
     main_parser.add_argument('--task_id', type=int, required=True, help=f'ID of the task to run. Options are: {"\n".join(f"{k}: {v}" for k, v in tasks.items())}', choices=list(tasks.keys()), default=0)
     main_parser.add_argument('--th', '--task_help', action="store_true", help="Help for the task specified by --task_id")
@@ -85,7 +84,7 @@ if __name__ == '__main__':
     ### If args has -h or --help, print help
     if any(x in remaining_args for x in ['-th', '--task_help']):
         task_id = args.task_id
-        hander, task_parser = tasks_handler_map[task_id]
+        task_handler, task_parser = tasks_handler_map[task_id]
         print("Help for task:", tasks[task_id])
         task_parser().print_help()
         exit(0)
@@ -94,6 +93,11 @@ if __name__ == '__main__':
 
 
     task_id = args.task_id
-    hander, task_parser = tasks_handler_map[task_id]
+    task_handler, task_parser = tasks_handler_map[task_id]
     task_args = task_parser().parse_args(remaining_args)
-    hander(task_args)
+    task_handler(task_args)
+    
+    
+if __name__ == '__main__':
+    main()
+