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

glam4cm/downstream_tasks/bert_node_classification.py

@@ -1,12 +1,19 @@
+import numpy as np
 from glam4cm.models.hf import get_model
-from glam4cm.downstream_tasks.common_args import get_bert_args_parser, get_common_args_parser, get_config_params
+from glam4cm.downstream_tasks.common_args import (
+    get_bert_args_parser, 
+    get_common_args_parser, 
+    get_config_params,
+    get_config_str
+)
 import os
 from transformers import TrainingArguments, Trainer
 from glam4cm.data_loading.graph_dataset import GraphNodeDataset
 from glam4cm.data_loading.utils import oversample_dataset
-from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.downstream_tasks.utils import get_logging_steps, get_models_dataset
+from glam4cm.settings import NODE_CLS_TASK, results_dir
 from glam4cm.tokenization.special_tokens import *
-
+from sklearn.model_selection import StratifiedKFold
 
 from sklearn.metrics import (
     accuracy_score, 
@@ -16,7 +23,7 @@ from sklearn.metrics import (
 )
 
 from glam4cm.tokenization.utils import get_tokenizer
-from glam4cm.utils import merge_argument_parsers, set_seed
+from glam4cm.utils import merge_argument_parsers, set_encoded_labels, set_seed
 
 
 
@@ -54,11 +61,24 @@ def get_parser():
     return parser
 
 
-
 def run(args):
     set_seed(args.seed)
+    dataset_name = args.dataset
+    print("Training model")
+    output_dir = os.path.join(
+        results_dir,
+        dataset_name,
+        f'LM_{NODE_CLS_TASK}',
+        f'{args.node_cls_label}',
+        get_config_str(args)
+    )
+
+    # if os.path.exists(output_dir):
+    #     print(f"Output directory {output_dir} already exists. Exiting.")
+    #     exit(0)
 
     config_params = dict(
+        include_dummies = args.include_dummies,
         min_enr = args.min_enr,
         min_edges = args.min_edges,
         remove_duplicates = args.remove_duplicates,
@@ -71,92 +91,110 @@ def run(args):
 
     print("Loaded dataset")
 
-    graph_data_params = get_config_params(args)
+    graph_data_params = {**get_config_params(args), 'task_type': NODE_CLS_TASK}
     print("Loading graph dataset")
-    graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
-    print("Loaded graph dataset")
-
-    assert hasattr(graph_dataset, f'num_nodes_{args.node_cls_label}'), f"Dataset does not have node_{args.node_cls_label} attribute"
-    num_labels = getattr(graph_dataset, f"num_nodes_{args.node_cls_label}")
-
-    model_name = args.model_name
-    tokenizer = get_tokenizer(model_name, use_special_tokens=args.use_special_tokens)
-
-    print("Getting node classification data")
-    bert_dataset = graph_dataset.get_node_classification_lm_data(
-        label=args.node_cls_label,
-        tokenizer=tokenizer,
-        distance=distance,
-    )
-
-    # exit(0)
-
-    if args.oversampling_ratio != -1:
-        ind_w_oversamples = oversample_dataset(bert_dataset['train'])
-        bert_dataset['train'].inputs = bert_dataset['train'][ind_w_oversamples]
-
-    model = get_model(
-        args.ckpt if args.ckpt else model_name, 
-        num_labels=2, 
-        len_tokenizer=len(tokenizer)
-    )
-
-    if args.freeze_pretrained_weights:
-        for param in model.base_model.parameters():
-            param.requires_grad = False
-
     
-    print("Training model")
-    output_dir = os.path.join(
-        'results',
-        dataset_name,
-        'node_cls',
-        f'{args.node_cls_label}',
-        f"{graph_dataset.config_hash}",
-    )
-
-    logs_dir = os.path.join(
-        'logs',
-        dataset_name,
-        'node_cls',
-        f'{args.node_cls_label}',
-        f"{graph_dataset.config_hash}",
-    )
-
-    print("Output Dir: ", output_dir)
-    print("Logs Dir: ", logs_dir)
-    print("Len Train Dataset: ", len(bert_dataset['train']))
-    print("Len Test Dataset: ", len(bert_dataset['test']))
-
-    training_args = TrainingArguments(
-        output_dir=output_dir,
-        num_train_epochs=args.num_epochs,
-        per_device_train_batch_size=args.train_batch_size,
-        per_device_eval_batch_size=args.eval_batch_size,
-        weight_decay=0.01,
-        logging_dir=logs_dir,
-        logging_steps=args.num_log_steps,
-        eval_strategy='steps',
-        eval_steps=args.num_eval_steps,
-        save_steps=args.num_save_steps,
-        save_total_limit=2,
-        load_best_model_at_end=True,
-        fp16=True,
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=bert_dataset['train'],
-        eval_dataset=bert_dataset['test'],
-        compute_metrics=compute_metrics
-    )
-
-    trainer.train()
-    results = trainer.evaluate()
-    print(results)
-
-    trainer.save_model()
+    k = int(1 / args.test_ratio)
+    
+    for i in range(k):
+        set_seed(np.random.randint(0, 1000))
+        graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
+        print("Loaded graph dataset")
+
+        assert hasattr(graph_dataset, f'num_nodes_{args.node_cls_label}'), f"Dataset does not have node_{args.node_cls_label} attribute"
+        num_labels = getattr(graph_dataset, f"num_nodes_{args.node_cls_label}")
+
+        model_name = args.model_name
+        tokenizer = get_tokenizer(model_name, use_special_tokens=args.use_special_tokens)
+
+        print("Getting node classification data")
+        bert_dataset = graph_dataset.get_node_classification_lm_data(
+            label=args.node_cls_label,
+            tokenizer=tokenizer,
+            distance=distance,
+        )
+
+        # exit(0)
+
+        if args.oversampling_ratio != -1:
+            ind_w_oversamples = oversample_dataset(bert_dataset['train'])
+            bert_dataset['train'].inputs = bert_dataset['train'][ind_w_oversamples]
+        
+        
+        model = get_model(
+            args.ckpt if args.ckpt else model_name, 
+            num_labels=num_labels, 
+            len_tokenizer=len(tokenizer), 
+            trust_remote_code=args.trust_remote_code
+        )
+
+        if args.freeze_pretrained_weights:
+            for param in model.base_model.parameters():
+                param.requires_grad = False
+
+
+        logs_dir = os.path.join(
+            'logs',
+            dataset_name,
+            f'BERT_{NODE_CLS_TASK}',
+            f'{args.node_cls_label}',
+            f"{graph_dataset.config_hash}_{i}",
+        )
+
+        print("Output Dir: ", output_dir)
+        print("Logs Dir: ", logs_dir)
+        print("Len Train Dataset: ", len(bert_dataset['train']))
+        print("Len Test Dataset: ", len(bert_dataset['test']))
+        
+        train_dataset = bert_dataset['train']
+        test_dataset = bert_dataset['test']
+        set_encoded_labels(train_dataset, test_dataset)
+
+        
+        print("Num epochs: ", args.num_epochs)
+        
+        logging_steps = get_logging_steps(
+            len(train_dataset), 
+            args.num_epochs, 
+            args.train_batch_size
+        )
+
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=args.num_epochs,
+            per_device_train_batch_size=args.train_batch_size,
+            per_device_eval_batch_size=args.eval_batch_size,
+            weight_decay=0.01,
+            logging_dir=logs_dir,
+            logging_steps=logging_steps,
+            eval_strategy='steps',
+            eval_steps=logging_steps,
+            # save_steps=args.num_save_steps,
+            # save_total_limit=2,
+            # load_best_model_at_end=True,
+            fp16=True,
+            save_strategy="no"
+        )
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=bert_dataset['train'],
+            eval_dataset=bert_dataset['test'],
+            compute_metrics=compute_metrics
+        )
+
+        trainer.train()
+        results = trainer.evaluate()
+        
+        # with open(os.path.join(output_dir, 'results.txt'), 'a') as f:
+        #     f.write(str(results))
+        #     f.write('\n')
+        
+        print(results)
+
+        trainer.save_model()
+        break
 
 
 if __name__ == '__main__':

glam4cm/downstream_tasks/cm_gpt_node_classification.py

@@ -1,6 +1,8 @@
 import os
 from glam4cm.downstream_tasks.common_args import (
-    get_common_args_parser, 
+    get_common_args_parser,
+    get_config_params,
+    get_config_str, 
     get_gpt_args_parser
 )
 
@@ -9,7 +11,8 @@ 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
+from glam4cm.utils import merge_argument_parsers, set_encoded_labels, set_seed
+from glam4cm.settings import NODE_CLS_TASK, results_dir
 
 
 def get_parser():
@@ -25,29 +28,68 @@ def get_parser():
 def run(args):
     set_seed(args.seed)
 
-    tokenizer = get_tokenizer('bert-base-cased', special_tokens=True)
+    tokenizer = get_tokenizer('bert-base-cased', use_special_tokens=args.use_special_tokens)
 
-    models_dataset_params = dict(
-        language='en',
+    set_seed(args.seed)
+    dataset_name = args.dataset
+    print("Training model")
+    output_dir = os.path.join(
+        results_dir,
+        dataset_name,
+        f'LM_{NODE_CLS_TASK}',
+        f'{args.node_cls_label}',
+        get_config_str(args)
     )
 
-    graph_params = dict(
-        use_special_tokens=args.use_special_tokens,
-        distance=args.distance,
-        reload = args.reload
+    # if os.path.exists(output_dir):
+    #     print(f"Output directory {output_dir} already exists. Exiting.")
+    #     exit(0)
+
+    config_params = dict(
+        include_dummies = args.include_dummies,
+        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
+    distance = args.distance
+    dataset = get_models_dataset(dataset_name, **config_params)
+
+    print("Loaded dataset")
 
-    models_dataset = get_models_dataset(args.dataset, **models_dataset_params)
-    graph_dataset = GraphNodeDataset(models_dataset, **graph_params)
+    graph_data_params = {**get_config_params(args), 'task_type': NODE_CLS_TASK}
+    print("Loading graph dataset")
+    graph_dataset = GraphNodeDataset(dataset, **graph_data_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,
+        distance=args.distance,
+    )
+    
+    set_encoded_labels(node_label_dataset['train'], node_label_dataset['test'])
+
+    print("Training model")
+    output_dir = os.path.join(
+        results_dir,
+        args.dataset,
+        f'LM_{NODE_CLS_TASK}',
+        f'{args.node_cls_label}',
+    )
+
+    logs_dir = os.path.join(
+        'logs',
+        args.dataset,
+        f'CMGPT_{NODE_CLS_TASK}',
+        f'{args.node_cls_label}',
+        f"{graph_dataset.config_hash}",
     )
 
+
     if args.pretr and os.path.exists(args.pretr):
         print(f"Loading pretrained model from {args.pretr}")
         cmgpt = CMGPT.from_pretrained(f"{args.pretr}")
@@ -60,7 +102,8 @@ def run(args):
             n_layer=args.n_layer,
             n_head=args.n_head,
         )
-    
+    print(f"Train dataset size: {len(node_label_dataset['train'])}")
+    print(f"Test dataset size: {len(node_label_dataset['test'])}")
     cmgpt_classifier = CMGPTClassifier(cmgpt, num_classes=getattr(graph_dataset, f"num_nodes_{args.node_cls_label}"))
 
     trainer = CMGPTTrainer(
@@ -68,9 +111,12 @@ def run(args):
         train_dataset=node_label_dataset['train'],
         test_dataset=node_label_dataset['test'],
         batch_size=args.batch_size, 
-        num_epochs=args.num_epochs
+        num_epochs=args.num_epochs,
+        log_dir=logs_dir,
+        results_dir=output_dir,
     )
 
     trainer.train()
 
-    trainer.save_model()
+    trainer.save_model()
+    

glam4cm/downstream_tasks/common_args.py

@@ -1,13 +1,35 @@
 from argparse import ArgumentParser
 from glam4cm.settings import (
+    MODERN_BERT,
     BERT_MODEL,
     WORD2VEC_MODEL,
     TFIDF_MODEL
 )
 
+def get_config_str(args):
+    config_str = ""
+    if args.use_attributes:
+        config_str += "_attrs"
+    if args.use_edge_label:
+        config_str += "_el"
+    if args.use_edge_types:
+        config_str += "_et"
+    if args.use_node_types:
+        config_str += "_nt"
+    if args.use_special_tokens:
+        config_str += "_st"
+    if args.no_labels:
+        config_str += "_nolb"
+    config_str += f"_{args.node_cls_label}" if args.node_cls_label else ""
+    config_str += f"_{args.edge_cls_label}" if args.edge_cls_label else ""
+    config_str += f"_{args.distance}" if args.distance else ""
+
+    return config_str
+
 
 def get_config_params(args):
     common_params = dict(
+        
         distance=args.distance,
         reload=args.reload,
         test_ratio=args.test_ratio,
@@ -17,6 +39,8 @@ def get_config_params(args):
         use_edge_types=args.use_edge_types,
         use_edge_label=args.use_edge_label,
         no_labels=args.no_labels,
+        
+        node_topk=args.node_topk,
 
         use_special_tokens=args.use_special_tokens,
 
@@ -58,6 +82,7 @@ def get_common_args_parser():
         ]
     )
     parser.add_argument('--remove_duplicates', action='store_true')
+    parser.add_argument('--include_dummies', 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)
@@ -72,6 +97,8 @@ def get_common_args_parser():
 
     parser.add_argument('--node_cls_label', type=str, default=None)
     parser.add_argument('--edge_cls_label', type=str, default=None)
+    
+    parser.add_argument('--node_topk', type=int, default=-1)
 
 
     parser.add_argument('--limit', type=int, default=-1)
@@ -84,9 +111,11 @@ def get_common_args_parser():
     parser.add_argument(
         '--embed_model_name', 
         type=str, 
-        default='bert-base-uncased', 
-        choices=[BERT_MODEL, WORD2VEC_MODEL, TFIDF_MODEL]
+        default=MODERN_BERT, 
+        choices=[MODERN_BERT, BERT_MODEL, WORD2VEC_MODEL, TFIDF_MODEL]
     )
+    
+    parser.add_argument('--trust_remote_code', action='store_true')
     parser.add_argument('--max_length', type=int, default=512)
     parser.add_argument('--ckpt', type=str, default=None)
     
@@ -134,7 +163,7 @@ 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('--model_name', type=str, default='answerdotai/ModernBERT-base')
 
     parser.add_argument('--warmup_steps', type=int, default=200)
     parser.add_argument('--num_log_steps', type=int, default=200)
@@ -148,7 +177,6 @@ def get_bert_args_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)

glam4cm/downstream_tasks/gnn_edge_classification.py

@@ -1,12 +1,16 @@
 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.settings import EDGE_CLS_TASK, results_dir
 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
+from glam4cm.utils import set_seed, merge_argument_parsers, 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():
@@ -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,
@@ -32,13 +37,21 @@ def run(args):
     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}
+    graph_data_params = {**graph_data_params, 'task_type': EDGE_CLS_TASK}
+    print("Using model: ", graph_data_params['embed_model_name'])
+    if args.ckpt:
+        print("Using checkpoint: ", args.ckpt)
+        
+    # if args.use_embeddings:
+    #     graph_data_params['embed_model_name'] = os.path.join(results_dir, dataset_name, f'{args.edge_cls_label}')
 
     print("Loading graph dataset")
     graph_dataset = GraphEdgeDataset(dataset, **graph_data_params)
     print("Loaded graph dataset")
 
     graph_torch_data = graph_dataset.get_torch_dataset()
+    exclude_labels = getattr(graph_dataset, f"edge_exclude_{args.edge_cls_label}")
+    set_torch_encoding_labels(graph_torch_data, f"edge_{args.edge_cls_label}", exclude_labels)
 
     input_dim = graph_torch_data[0].x.shape[1]
 
@@ -59,12 +72,16 @@ def run(args):
 
     edge_dim = graph_dataset[0].data.edge_attr.shape[1] if args.use_edge_attrs else None
 
+    ue = "" if not args.use_edge_attrs else "_ue"
+    
     logs_dir = os.path.join(
         "logs",
         dataset_name,
-        "gnn_edge_cls",
+        f"GNN_{EDGE_CLS_TASK}",
+        f"{args.edge_cls_label}{ue}",
         f"{graph_dataset.config_hash}",
     )
+    
 
     gnn_conv_model = GNNConv(
         model_name=model_name,
@@ -99,8 +116,8 @@ def run(args):
         num_epochs=args.num_epochs,
         batch_size=args.batch_size,
         use_edge_attrs=args.use_edge_attrs,
-        logs_dir=logs_dir
+        logs_dir=logs_dir,
     )
 
     print("Training GNN Edge Classification model")
-    trainer.run()
+    trainer.run()

glam4cm/downstream_tasks/gnn_graph_cls.py

@@ -1,6 +1,7 @@
 import os
 from glam4cm.data_loading.graph_dataset import GraphNodeDataset
 from glam4cm.models.gnn_layers import GNNConv, GraphClassifer
+from glam4cm.settings import GRAPH_CLS_TASK, DUMMY_GRAPH_CLS_TASK, results_dir
 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
@@ -21,6 +22,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,
@@ -30,8 +32,10 @@ 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': GRAPH_CLS_TASK if not args.include_dummies else DUMMY_GRAPH_CLS_TASK}
+    # if args.use_embeddings:
+    #     graph_data_params['ckpt'] = os.path.join(results_dir, dataset_name, f'{args.cls_label}')
 
     print("Loading graph dataset")
     graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
@@ -40,8 +44,16 @@ def run(args):
     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}")
+    
+    if args.include_dummies:
+        import numpy as np
+        dummy_class = int(graph_dataset.graph_label_map_label.transform(['dummy'])[0])
+        for g, l in zip(graph_dataset, [int(g.data.graph_label[0]) == dummy_class for g in graph_dataset]):
+            setattr(g.data, f"graph_{args.cls_label}", np.array([int(l)]))
+        num_classes = 2
+        
+    
     model_name = args.gnn_conv_model
     hidden_dim = args.hidden_dim
     output_dim = args.output_dim
@@ -53,14 +65,15 @@ def run(args):
     aggregation = args.aggregation
 
     input_dim = graph_dataset[0].data.x.shape[1]
-
+    ue = "" if not args.use_edge_attrs else "_ue"
     logs_dir = os.path.join(
         "logs",
         dataset_name,
-        "gnn_graph_cls",
+        f"GNN_{GRAPH_CLS_TASK}{ue}",
         f"{graph_dataset.config_hash}",
     )
 
+    fold_id = 0
     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
@@ -94,7 +107,7 @@ def run(args):
             num_epochs=args.num_epochs,
             batch_size=args.batch_size,
             use_edge_attrs=args.use_edge_attrs,
-            logs_dir=logs_dir
+            logs_dir=logs_dir + f"_{fold_id}",
         )
 
         trainer.run()