glam4cm 0.1.0__py3-none-any.whl

glam4cm/downstream_tasks/cm_gpt_node_classification.py

@@ -0,0 +1,76 @@
+import os
+from glam4cm.downstream_tasks.common_args import (
+    get_common_args_parser, 
+    get_gpt_args_parser
+)
+
+from glam4cm.data_loading.graph_dataset import GraphNodeDataset
+from glam4cm.models.cmgpt import CMGPT, CMGPTClassifier
+from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.tokenization.utils import get_tokenizer
+from glam4cm.trainers.cm_gpt_trainer import CMGPTTrainer
+from glam4cm.utils import merge_argument_parsers, set_seed
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    bert_parser = get_gpt_args_parser()
+    parser = merge_argument_parsers(common_parser, bert_parser)
+
+    parser.add_argument('--cls_label', type=str)
+    parser.add_argument('--pretr', type=str, default=None)
+    return parser
+
+
+def run(args):
+    set_seed(args.seed)
+
+    tokenizer = get_tokenizer('bert-base-cased', special_tokens=True)
+
+    models_dataset_params = dict(
+        language='en',
+    )
+
+    graph_params = dict(
+        use_special_tokens=args.use_special_tokens,
+        distance=args.distance,
+        reload = args.reload
+    )
+
+    models_dataset = get_models_dataset(args.dataset, **models_dataset_params)
+    graph_dataset = GraphNodeDataset(models_dataset, **graph_params)
+
+    assert hasattr(graph_dataset, f'num_nodes_{args.node_cls_label}'), f"Dataset does not have node labels for {args.node_cls_label}"
+
+    node_label_dataset = graph_dataset.get_node_classification_lm_data(
+        args.node_cls_label,
+        tokenizer=tokenizer,
+        distance=1,
+    )
+
+    if args.pretr and os.path.exists(args.pretr):
+        print(f"Loading pretrained model from {args.pretr}")
+        cmgpt = CMGPT.from_pretrained(f"{args.pretr}")
+    else:
+        print("Creating new model")
+        cmgpt = CMGPT(
+            vocab_size=len(tokenizer),
+            embed_dim=args.embed_dim,
+            block_size=args.block_size,
+            n_layer=args.n_layer,
+            n_head=args.n_head,
+        )
+    
+    cmgpt_classifier = CMGPTClassifier(cmgpt, num_classes=getattr(graph_dataset, f"num_nodes_{args.node_cls_label}"))
+
+    trainer = CMGPTTrainer(
+        cmgpt_classifier, 
+        train_dataset=node_label_dataset['train'],
+        test_dataset=node_label_dataset['test'],
+        batch_size=args.batch_size, 
+        num_epochs=args.num_epochs
+    )
+
+    trainer.train()
+
+    trainer.save_model()

glam4cm/downstream_tasks/cm_gpt_pretraining.py

@@ -0,0 +1,64 @@
+from sklearn.model_selection import train_test_split
+from glam4cm.downstream_tasks.common_args import (
+    get_common_args_parser, 
+    get_gpt_args_parser
+)
+
+from glam4cm.data_loading.graph_dataset import get_models_gpt_dataset
+from glam4cm.models.cmgpt import CMGPT
+from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.tokenization.utils import get_tokenizer
+from glam4cm.trainers.cm_gpt_trainer import CMGPTTrainer
+from glam4cm.utils import merge_argument_parsers, set_seed
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    bert_parser = get_gpt_args_parser()
+    parser = merge_argument_parsers(common_parser, bert_parser)
+    return parser
+
+
+def run(args):
+    set_seed(args.seed)
+
+    tokenizer = get_tokenizer('bert-base-cased', special_tokens=True)
+
+    models_dataset_params = dict(
+        language='en',
+    )
+
+    graph_params = dict(
+        use_special_tokens=True,
+        distance=1,
+    )
+
+    models_dataset = get_models_dataset(args.dataset, **models_dataset_params)
+    graph_dataset = get_models_gpt_dataset(
+        models_dataset, 
+        tokenizer, 
+        **graph_params
+    )
+
+    train_dataset, test_dataset = train_test_split(graph_dataset, test_size=0.05)
+
+    cmgpt = CMGPT(
+        vocab_size=len(tokenizer),
+        embed_dim=args.embed_dim,
+        block_size=args.block_size,
+        n_layer=args.n_layer,
+        n_head=args.n_head,
+    )
+
+    trainer = CMGPTTrainer(
+        cmgpt, 
+        train_dataset=train_dataset,
+        test_dataset=test_dataset,
+        test_ratio=0.05, 
+        batch_size=args.batch_size, 
+        num_epochs=args.num_epochs
+    )
+
+    trainer.train()
+
+    trainer.save_model()

glam4cm/downstream_tasks/common_args.py

@@ -0,0 +1,160 @@
+from argparse import ArgumentParser
+from glam4cm.settings import (
+    BERT_MODEL,
+    WORD2VEC_MODEL,
+    TFIDF_MODEL
+)
+
+
+def get_config_params(args):
+    common_params = dict(
+        distance=args.distance,
+        reload=args.reload,
+        test_ratio=args.test_ratio,
+
+        use_attributes=args.use_attributes,
+        use_node_types=args.use_node_types,
+        use_edge_types=args.use_edge_types,
+        use_edge_label=args.use_edge_label,
+        no_labels=args.no_labels,
+
+        use_special_tokens=args.use_special_tokens,
+
+        use_embeddings=args.use_embeddings,
+        embed_model_name=args.embed_model_name,
+        ckpt=args.ckpt,
+        
+        no_shuffle=args.no_shuffle,
+        randomize_ne=args.randomize_ne,
+        randomize_ee=args.randomize_ee,
+        random_embed_dim=args.random_embed_dim,
+
+        limit = args.limit,
+
+        node_cls_label=args.node_cls_label,
+        edge_cls_label=args.edge_cls_label,
+    )
+    
+
+    return common_params
+
+
+def get_common_args_parser():
+    parser = ArgumentParser()
+
+    parser.add_argument('--seed', type=int, default=42)
+
+    ### Models Dataset Creation
+    parser.add_argument(
+        '--dataset', 
+        type=str, 
+        default='ecore_555', 
+        choices=[
+            'modelset', 
+            'ecore_555', 
+            'mar-ecore-github', 
+            'eamodelset',
+            'ontouml'
+        ]
+    )
+    parser.add_argument('--remove_duplicates', action='store_true')
+    parser.add_argument('--reload', action='store_true')
+    parser.add_argument('--min_enr', type=float, default=-1.0)
+    parser.add_argument('--min_edges', type=int, default=-1)
+    parser.add_argument('--language', type=str, default='en')
+    
+    parser.add_argument('--use_attributes', action='store_true')
+    parser.add_argument('--use_edge_label', action='store_true')
+    parser.add_argument('--use_edge_types', action='store_true')
+    parser.add_argument('--use_node_types', action='store_true')
+    parser.add_argument('--use_special_tokens', action='store_true')
+    parser.add_argument('--no_labels', action='store_true')
+
+    parser.add_argument('--node_cls_label', type=str, default=None)
+    parser.add_argument('--edge_cls_label', type=str, default=None)
+
+
+    parser.add_argument('--limit', type=int, default=-1)
+
+
+    ### Model Dataset Loading
+    parser.add_argument('--distance', type=int, default=0)
+    parser.add_argument('--use_embeddings', action='store_true')
+    parser.add_argument('--regen_embeddings', action='store_true')
+    parser.add_argument(
+        '--embed_model_name', 
+        type=str, 
+        default='bert-base-uncased', 
+        choices=[BERT_MODEL, WORD2VEC_MODEL, TFIDF_MODEL]
+    )
+    parser.add_argument('--max_length', type=int, default=512)
+    parser.add_argument('--ckpt', type=str, default=None)
+    
+
+    parser.add_argument('--no_shuffle', action='store_true')
+    parser.add_argument('--randomize_ne', action='store_true')
+    parser.add_argument('--randomize_ee', action='store_true')
+    parser.add_argument('--random_embed_dim', type=int, default=128)
+
+
+    parser.add_argument('--test_ratio', type=float, default=0.2)
+    parser.add_argument('--add_negative_train_samples', action='store_true')
+    parser.add_argument('--neg_sampling_ratio', type=int, default=1)
+
+    parser.add_argument('--num_epochs', type=int, default=100)
+    parser.add_argument('--batch_size', type=int, default=32)
+
+    return parser
+
+
+def get_gnn_args_parser():
+    parser = ArgumentParser()
+    parser.add_argument('--num_conv_layers', type=int, default=3)
+    parser.add_argument('--num_mlp_layers', type=int, default=3)
+    parser.add_argument('--num_heads', type=int, default=None)
+
+    parser.add_argument('--input_dim', type=int, default=768)
+    parser.add_argument('--hidden_dim', type=int, default=128)
+    parser.add_argument('--output_dim', type=int, default=128)
+
+    parser.add_argument('--residual', action='store_true')
+    parser.add_argument('--bias', action='store_true')
+    parser.add_argument('--l_norm', action='store_true')
+    parser.add_argument('--aggregation', type=str, default='sum')
+    parser.add_argument('--dropout', type=float, default=0.3)
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--gnn_conv_model', type=str, default='SAGEConv')
+
+    parser.add_argument('--use_edge_attrs', action='store_true')
+    
+    return parser
+
+
+def get_bert_args_parser():
+    parser = ArgumentParser()
+
+    parser.add_argument('--freeze_pretrained_weights', action='store_true')
+    parser.add_argument('--model_name', type=str, default='bert-base-uncased')
+
+    parser.add_argument('--warmup_steps', type=int, default=200)
+    parser.add_argument('--num_log_steps', type=int, default=200)
+    parser.add_argument('--num_eval_steps', type=int, default=200)
+    parser.add_argument('--num_save_steps', type=int, default=200)
+    parser.add_argument('--train_batch_size', type=int, default=2)
+    parser.add_argument('--eval_batch_size', type=int, default=128)
+    return parser
+
+
+def get_gpt_args_parser():
+    parser = ArgumentParser()
+    parser.add_argument('--model_name', type=str, default='gpt2')
+    parser.add_argument('--use_special_tokens', action='store_true')
+
+    parser.add_argument('--warmup_steps', type=int, default=200)
+    parser.add_argument('--blocks', type=int, default=6)
+    parser.add_argument('--block_size', type=int, default=8)
+    parser.add_argument('--n_head', type=int, default=8)
+    parser.add_argument('--embed_dim', type=int, default=768)
+    parser.add_argument('--n_layer', type=int, default=6)
+    parser.add_argument('--lr', type=float, default=1e-5)
+    return parser

glam4cm/downstream_tasks/create_dataset.py

@@ -0,0 +1,51 @@
+import random
+import torch
+import numpy as np
+
+from glam4cm.data_loading.models_dataset import ArchiMateDataset, EcoreDataset
+from glam4cm.data_loading.graph_dataset import (
+    GraphNodeDataset,
+    GraphEdgeDataset
+)
+from glam4cm.downstream_tasks.common_args import get_common_args_parser
+
+
+def get_parser():
+    parser = get_common_args_parser()
+    return parser.parse_args()
+
+
+def run(args):
+    
+    seed = 42
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+    
+    config_params = dict(
+        min_enr = args.min_enr,
+        min_edges = args.min_edges,
+    )
+    ecore = EcoreDataset('ecore_555', reload=args.reload, **config_params)
+    modelset = EcoreDataset('modelset', reload=args.reload, remove_duplicates=True, **config_params)
+    mar = EcoreDataset('mar-ecore-github', reload=args.reload, **config_params)
+    eamodelset = ArchiMateDataset('eamodelset', reload=args.reload, **config_params)
+
+    graph_data_params = dict(
+        distance=args.distance,
+        add_negative_train_samples=args.add_neg_samples,
+        neg_sampling_ratio=1,
+    )
+
+    GraphEdgeDataset(ecore, reload=False, **graph_data_params)
+    GraphEdgeDataset(modelset, reload=True, **graph_data_params)
+    GraphEdgeDataset(mar, reload=True, **graph_data_params)
+    GraphEdgeDataset(eamodelset, reload=True, **graph_data_params)
+
+
+    GraphNodeDataset(ecore, reload=False, **graph_data_params)
+    GraphNodeDataset(modelset, reload=True, **graph_data_params)
+    GraphNodeDataset(mar, reload=True, **graph_data_params)
+    GraphNodeDataset(eamodelset, reload=True, **graph_data_params)

glam4cm/downstream_tasks/gnn_edge_classification.py

@@ -0,0 +1,106 @@
+import os
+from glam4cm.data_loading.graph_dataset import GraphEdgeDataset
+from glam4cm.models.gnn_layers import GNNConv, EdgeClassifer
+from glam4cm.settings import LP_TASK_EDGE_CLS
+from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.tokenization.special_tokens import *
+from glam4cm.trainers.gnn_edge_classifier import GNNEdgeClassificationTrainer as Trainer
+from glam4cm.utils import set_seed, merge_argument_parsers
+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()
+    parser = merge_argument_parsers(common_parser, gnn_parser)
+    return parser
+
+
+def run(args):
+
+    set_seed(args.seed)
+    
+    config_params = dict(
+        min_enr = args.min_enr,
+        min_edges = args.min_edges,
+        remove_duplicates = args.remove_duplicates,
+        reload = args.reload,
+        language = args.language
+    )
+    dataset_name = args.dataset
+
+    dataset = get_models_dataset(dataset_name, **config_params)
+
+    graph_data_params = get_config_params(args)
+    graph_data_params = {**graph_data_params, 'task_type': LP_TASK_EDGE_CLS}
+
+    print("Loading graph dataset")
+    graph_dataset = GraphEdgeDataset(dataset, **graph_data_params)
+    print("Loaded graph dataset")
+
+    graph_torch_data = graph_dataset.get_torch_dataset()
+
+    input_dim = graph_torch_data[0].x.shape[1]
+
+    model_name = args.gnn_conv_model
+    hidden_dim = args.hidden_dim
+    output_dim = args.output_dim
+    num_conv_layers = args.num_conv_layers
+    num_mlp_layers = args.num_mlp_layers
+    num_heads = args.num_heads
+    residual = False
+    l_norm = args.l_norm
+    dropout = args.dropout
+    aggregation = args.aggregation
+
+    num_edges_label = f"num_edges_{args.edge_cls_label}"
+    assert hasattr(graph_dataset, num_edges_label), f"Graph dataset does not have attribute {num_edges_label}"
+    num_classes = getattr(graph_dataset, num_edges_label)
+
+    edge_dim = graph_dataset[0].data.edge_attr.shape[1] if args.use_edge_attrs else None
+
+    logs_dir = os.path.join(
+        "logs",
+        dataset_name,
+        "gnn_edge_cls",
+        f"{graph_dataset.config_hash}",
+    )
+
+    gnn_conv_model = GNNConv(
+        model_name=model_name,
+        input_dim=input_dim,
+        hidden_dim=hidden_dim,
+        out_dim=output_dim,
+        num_layers=num_conv_layers,
+        num_heads=num_heads,
+        residual=residual,
+        l_norm=l_norm,
+        dropout=dropout,
+        aggregation=aggregation,
+        edge_dim=edge_dim
+    )
+
+    clf_input_dim = output_dim*num_heads if args.num_heads else output_dim
+    mlp_predictor = EdgeClassifer(
+        input_dim=clf_input_dim,
+        hidden_dim=hidden_dim,
+        num_layers=num_mlp_layers, 
+        num_classes=num_classes,
+        edge_dim=edge_dim,
+        bias=args.bias,
+    )
+
+    trainer = Trainer(
+        gnn_conv_model, 
+        mlp_predictor, 
+        graph_torch_data,
+        cls_label=args.edge_cls_label,
+        lr=args.lr,
+        num_epochs=args.num_epochs,
+        batch_size=args.batch_size,
+        use_edge_attrs=args.use_edge_attrs,
+        logs_dir=logs_dir
+    )
+
+    print("Training GNN Edge Classification model")
+    trainer.run()

glam4cm/downstream_tasks/gnn_graph_cls.py

@@ -0,0 +1,101 @@
+import os
+from glam4cm.data_loading.graph_dataset import GraphNodeDataset
+from glam4cm.models.gnn_layers import GNNConv, GraphClassifer
+from glam4cm.trainers.gnn_graph_classifier import GNNGraphClassificationTrainer as Trainer
+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
+from glam4cm.downstream_tasks.utils import get_models_dataset
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    gnn_parser = get_gnn_args_parser()
+    parser = merge_argument_parsers(common_parser, gnn_parser)
+
+    parser.add_argument('--cls_label', type=str, default='label')
+    parser.add_argument('--global_pool', type=str, default='mean')
+    return parser
+
+
+def run(args):
+    set_seed(args.seed)
+    
+    config_params = dict(
+        min_enr = args.min_enr,
+        min_edges = args.min_edges,
+        remove_duplicates = args.remove_duplicates,
+        reload = args.reload,
+        language = args.language
+    )
+    dataset_name = args.dataset
+
+    dataset = get_models_dataset(dataset_name, **config_params)
+
+    graph_data_params = get_config_params(args)
+
+    print("Loading graph dataset")
+    graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
+    print("Loaded graph dataset")
+
+    cls_label = f"num_graph_{args.cls_label}"
+    assert hasattr(graph_dataset, cls_label), f"Dataset does not have attribute {cls_label}"
+    num_classes = getattr(graph_dataset, cls_label)
+
+    print(f"Number of classes: {num_classes}")
+    model_name = args.gnn_conv_model
+    hidden_dim = args.hidden_dim
+    output_dim = args.output_dim
+    num_conv_layers = args.num_conv_layers
+    num_heads = args.num_heads
+    residual = True
+    l_norm = False
+    dropout = args.dropout
+    aggregation = args.aggregation
+
+    input_dim = graph_dataset[0].data.x.shape[1]
+
+    logs_dir = os.path.join(
+        "logs",
+        dataset_name,
+        "gnn_graph_cls",
+        f"{graph_dataset.config_hash}",
+    )
+
+    for datasets in graph_dataset.get_kfold_gnn_graph_classification_data():
+
+        edge_dim = graph_dataset[0].data.edge_attr.shape[1] if args.num_heads else None
+
+        gnn_conv_model = GNNConv(
+            model_name=model_name,
+            input_dim=input_dim,
+            hidden_dim=hidden_dim,
+            out_dim=output_dim,
+            num_layers=num_conv_layers,
+            num_heads=num_heads,
+            residual=residual,
+            l_norm=l_norm,
+            dropout=dropout,
+            aggregation=aggregation,
+            edge_dim=edge_dim,
+        )
+
+        clf_input_dim = output_dim*num_heads if args.num_heads else output_dim
+        classifier = GraphClassifer(
+            input_dim=clf_input_dim,
+            num_classes=num_classes,
+            global_pool=args.global_pool,
+        )
+
+        trainer = Trainer(
+            gnn_conv_model,
+            classifier, 
+            datasets,
+            lr=args.lr,
+            num_epochs=args.num_epochs,
+            batch_size=args.batch_size,
+            use_edge_attrs=args.use_edge_attrs,
+            logs_dir=logs_dir
+        )
+
+        trainer.run()
+        break

glam4cm/downstream_tasks/gnn_link_prediction.py

@@ -0,0 +1,109 @@
+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.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
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    gnn_parser = get_gnn_args_parser()
+    parser = merge_argument_parsers(common_parser, gnn_parser)
+    return parser
+
+
+def run(args):
+
+    set_seed(args.seed)
+    
+    config_params = dict(
+        min_enr = args.min_enr,
+        min_edges = args.min_edges,
+        remove_duplicates = args.remove_duplicates,
+        reload = args.reload,
+        language = args.language
+    )
+    dataset_name = args.dataset
+    dataset = get_models_dataset(dataset_name, **config_params)
+    
+    model_name = args.gnn_conv_model
+    hidden_dim = args.hidden_dim
+    output_dim = args.output_dim
+    num_conv_layers = args.num_conv_layers
+    num_mlp_layers = args.num_mlp_layers
+    num_heads = args.num_heads
+    residual = True
+    l_norm = args.l_norm
+    dropout = args.dropout
+    aggregation = args.aggregation
+
+    graph_data_params = get_config_params(args)
+    print("Loading graph dataset")
+    graph_dataset = GraphEdgeDataset(
+        dataset, 
+        dict(
+            **graph_data_params, 
+            add_negative_train_samples=args.add_negative_train_samples, 
+            neg_sampling_ratio=args.neg_sampling_ratio,
+            task=LP_TASK_LINK_PRED
+    ))
+
+    input_dim = graph_dataset[0].data.x.shape[1]
+
+    edge_dim = None
+    if args.use_edge_attrs:
+        if args.use_embeddings:
+            edge_dim = graph_dataset.embedder.embedding_dim
+        else:
+            edge_dim = graph_dataset[0].data.edge_attr.shape[1]
+    
+    gnn_conv_model = GNNConv(
+        model_name=model_name,
+        input_dim=input_dim,
+        hidden_dim=hidden_dim,
+        out_dim=output_dim,
+        num_layers=num_conv_layers,
+        num_heads=num_heads,
+        residual=residual,
+        l_norm=l_norm,
+        dropout=dropout,
+        aggregation=aggregation,
+        edge_dim=edge_dim
+    )
+
+    logs_dir = os.path.join(
+        "logs",
+        dataset_name,
+        "gnn_lp",
+        f'{graph_dataset.config_hash}',
+    )
+
+    clf_input_dim = gnn_conv_model.out_dim*num_heads if args.num_heads else output_dim
+    mlp_predictor = EdgeClassifer(
+        input_dim=clf_input_dim,
+        hidden_dim=hidden_dim,
+        num_layers=num_mlp_layers, 
+        num_classes=2,
+        edge_dim=edge_dim,
+        bias=False,
+    )
+
+    
+    trainer = Trainer(
+        gnn_conv_model, 
+        mlp_predictor, 
+        graph_dataset.get_torch_dataset(),
+        lr=args.lr,
+        num_epochs=args.num_epochs,
+        batch_size=args.batch_size,
+        use_edge_attrs=args.use_edge_attrs,
+        logs_dir=logs_dir
+    )
+
+
+    print("Training GNN Link Prediction model")
+    trainer.run()