glam4cm 0.1.0__py3-none-any.whl

glam4cm/downstream_tasks/bert_graph_classification.py

@@ -0,0 +1,137 @@
+import os
+from sklearn.metrics import (
+    accuracy_score, 
+    balanced_accuracy_score, 
+    f1_score, 
+    recall_score
+)
+from transformers import (
+    Trainer,
+    TrainingArguments
+)
+
+from glam4cm.data_loading.graph_dataset import GraphNodeDataset
+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.utils import get_models_dataset
+from glam4cm.tokenization.utils import get_tokenizer
+from glam4cm.utils import merge_argument_parsers, set_seed
+
+
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+    acc = (preds == labels).mean()
+    f1_macro = f1_score(labels, preds, average='macro')
+    accuracy = accuracy_score(labels, preds)
+    recall = recall_score(labels, preds, average='macro')
+    balanced_acc = balanced_accuracy_score(labels, preds)
+
+    return {
+        'balanced_accuracy': balanced_acc,
+        'accuracy': acc,
+        'f1_macro': f1_macro,
+        'precision': accuracy,
+        'recall': recall
+    }
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    bert_parser = get_bert_args_parser()
+    parser = merge_argument_parsers(common_parser, bert_parser)
+
+    parser.add_argument('--cls_label', type=str, default='label')
+    parser.add_argument('--remove_duplicate_graphs', action='store_true')
+    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")
+
+    model_name = args.model_name
+    tokenizer = get_tokenizer(model_name, args.use_special_tokens)
+
+    fold_id = 0
+    for classification_dataset in graph_dataset.get_kfold_lm_graph_classification_data(
+        tokenizer,
+        remove_duplicates=args.remove_duplicate_graphs
+    ):
+        train_dataset = classification_dataset['train']
+        test_dataset = classification_dataset['test']
+        num_labels = classification_dataset['num_classes']
+
+        print(len(train_dataset), len(test_dataset), num_labels)
+
+        print("Training model")
+        output_dir = os.path.join(
+            'results',
+            dataset_name,
+            f'graph_cls_',
+            f"{graph_dataset.config_hash}",
+        )
+
+        logs_dir = os.path.join(
+            'logs',
+            dataset_name,
+            f'graph_cls_',
+            f"{graph_dataset.config_hash}"
+        )
+
+        model = get_model(args.ckpt if args.ckpt else model_name, num_labels, len(tokenizer))
+
+        if args.freeze_pretrained_weights:
+            for param in model.base_model.parameters():
+                param.requires_grad = False
+
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=args.num_epochs,
+            eval_strategy="steps",
+            per_device_train_batch_size=args.train_batch_size,
+            per_device_eval_batch_size=args.eval_batch_size,
+            warmup_steps=200,
+            weight_decay=0.01,
+            learning_rate=5e-5,
+            logging_dir=logs_dir,
+            logging_steps=args.num_log_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 = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
+            compute_metrics=compute_metrics            
+        )
+
+        # Train the model
+        trainer.train()
+        results = trainer.evaluate()
+        print(results)
+        
+        fold_id += 1
+        break

glam4cm/downstream_tasks/bert_graph_classification_comp.py

@@ -0,0 +1,156 @@
+import os
+import json
+from argparse import ArgumentParser
+from random import shuffle
+
+import numpy as np
+from sklearn.model_selection import StratifiedKFold
+
+from sklearn.metrics import (
+    accuracy_score, 
+    balanced_accuracy_score, 
+    f1_score, 
+    recall_score
+)
+from transformers import (
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments
+)
+
+from glam4cm.data_loading.encoding import EncodingDataset
+from glam4cm.models.hf import get_model
+
+
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = np.argmax(pred.predictions, axis=1)
+    acc = (preds == labels).mean()
+    f1_macro = f1_score(labels, preds, average='macro')
+    accuracy = accuracy_score(labels, preds)
+    recall = recall_score(labels, preds, average='macro')
+    balanced_acc = balanced_accuracy_score(labels, preds)
+
+    return {
+        'balanced_accuracy': balanced_acc,
+        'accuracy': acc,
+        'f1_macro': f1_macro,
+        'precision': accuracy,
+        'recall': recall
+    }
+
+
+def get_parser():
+    parser = ArgumentParser()
+    parser.add_argument('--seed', type=int, default=42)
+    parser.add_argument('--dataset_name', type=str, default='ecore_555')
+    parser.add_argument('--model_name', type=str, default='bert-base-uncased')
+    parser.add_argument('--ckpt', type=str, default=None)
+    parser.add_argument('--max_length', type=int, default=512)
+    parser.add_argument('--k', type=int, default=10)
+
+    parser.add_argument('--num_epochs', type=int, default=10)
+
+    parser.add_argument('--warmup_steps', type=int, default=500)
+    parser.add_argument('--num_log_steps', type=int, default=50)
+    parser.add_argument('--num_eval_steps', type=int, default=50)
+    parser.add_argument('--num_save_steps', type=int, default=50)
+    parser.add_argument('--train_batch_size', type=int, default=2)
+    parser.add_argument('--eval_batch_size', type=int, default=128)
+    parser.add_argument('--lr', type=float, default=1e-5)
+
+    return parser
+
+
+def run(args):
+    dataset_name = args.dataset_name
+    model_name = args.model_name
+
+
+    texts = [
+        (g['txt'], g['labels'])
+        for file_name in os.listdir(f'datasets/{dataset_name}')
+        for g in json.load(open(f'datasets/{dataset_name}/{file_name}'))
+        if 'ecore' in file_name and file_name.endswith('.jsonl')
+    ]
+    shuffle(texts)
+    labels = [y for _, y in texts]
+    y_map = {label: i for i, label in enumerate(set(y for y in labels))}
+    y = [y_map[y] for y in labels]
+    n = len(texts)
+
+    texts = [text for text, _ in texts]
+
+    num_labels = len(y_map)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    k = args.k
+    kfold = StratifiedKFold(n_splits=k, shuffle=True, random_state=args.seed)
+
+    i = 0
+    for train_idx, test_idx in kfold.split(np.zeros(n), np.zeros(n)):
+
+        print(f'Fold {i+1}/{k}')
+
+        train_texts = [texts[i] for i in train_idx]
+        test_texts = [texts[i] for i in test_idx]
+        train_y = [y[i] for i in train_idx]
+        test_y = [y[i] for i in test_idx]
+
+
+        print(f'Train: {len(train_texts)}, Test: {len(test_texts)}', num_labels)
+
+        train_dataset = EncodingDataset(tokenizer, train_texts, train_y)
+        test_dataset = EncodingDataset(tokenizer, test_texts, test_y)
+
+        model = get_model(args.ckpt if args.ckpt else model_name, num_labels, len(tokenizer))
+
+        print("Training model")
+        output_dir = os.path.join(
+            'results',
+            dataset_name,
+            'graph_cls_comp',
+        )
+
+        logs_dir = os.path.join(
+            'logs',
+            dataset_name,
+            'graph_cls_comp',
+        )
+
+        print("Running epochs: ", args.num_epochs)
+
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            num_train_epochs=args.num_epochs,
+            eval_strategy="steps",
+            per_device_train_batch_size=args.train_batch_size,
+            per_device_eval_batch_size=args.eval_batch_size,
+            warmup_steps=500,
+            weight_decay=0.01,
+            logging_dir=logs_dir,
+            logging_steps=10,
+            eval_steps=10,
+            save_total_limit=2,
+            load_best_model_at_end=True,
+            fp16=True
+        )
+
+        # Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=test_dataset,
+            compute_metrics=compute_metrics            
+        )
+
+        # Train the model
+        trainer.train()
+        results = trainer.evaluate()
+        print(results)
+
+        i += 1
+        

glam4cm/downstream_tasks/bert_link_prediction.py

@@ -0,0 +1,145 @@
+from collections import Counter
+import os
+from transformers import TrainingArguments, Trainer
+from glam4cm.data_loading.graph_dataset import GraphEdgeDataset
+from glam4cm.models.hf import get_model
+from glam4cm.settings import LP_TASK_LINK_PRED
+from glam4cm.downstream_tasks.common_args import get_bert_args_parser, get_common_args_parser, get_config_params
+from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.tokenization.special_tokens import *
+
+
+from sklearn.metrics import (
+    f1_score, 
+    precision_score, 
+    balanced_accuracy_score,
+    recall_score
+)
+
+from glam4cm.tokenization.utils import get_tokenizer
+from glam4cm.utils import merge_argument_parsers, set_seed
+
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+    print(Counter(preds), Counter(labels))
+    acc = (preds == labels).mean()
+    # roc = roc_auc_score(labels, preds)
+    f1_macro = f1_score(labels, preds)
+    f1_micro = f1_score(labels, preds, )
+    precision = precision_score(labels, preds)
+    balanced_accuracy = balanced_accuracy_score(labels, preds)
+    recall = recall_score(labels, preds)
+
+    return {
+        'accuracy': acc,
+        'f1_macro': f1_macro,
+        'f1_micro': f1_micro,
+        'precision': precision,
+        'recall': recall,
+        'balanced_accuracy': balanced_accuracy
+    }
+
+
+
+def get_parser():
+
+    common_parser = get_common_args_parser()
+    bert_parser = get_bert_args_parser()
+    parser = merge_argument_parsers(common_parser, bert_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)
+
+    print("Loaded dataset")
+
+
+    graph_data_params = get_config_params(args)
+    graph_data_params = {**graph_data_params, 'task': LP_TASK_LINK_PRED}
+
+    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
+    ))
+    print("Loaded graph dataset")
+
+
+
+    model_name = args.model_name
+    tokenizer = get_tokenizer(model_name, args.use_special_tokens)
+
+
+    print("Getting link prediction data")
+    bert_dataset = graph_dataset.get_link_prediction_lm_data(
+        tokenizer=tokenizer,
+        task_type=LP_TASK_LINK_PRED
+    )
+
+    print("Training model")
+    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
+
+
+    output_dir = os.path.join(
+        'results',
+        dataset_name,
+        'lp',
+        f"{graph_dataset.config_hash}",
+    )
+
+    logs_dir = os.path.join(
+        'logs',
+        dataset_name,
+        'lp',
+        f"{graph_dataset.config_hash}",
+    )
+    
+    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=200,
+        eval_strategy='steps',
+        eval_steps=200,
+        save_steps=200,
+        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()
+    print(trainer.evaluate())
+    trainer.save_model()

glam4cm/downstream_tasks/bert_node_classification.py

@@ -0,0 +1,164 @@
+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
+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.tokenization.special_tokens import *
+
+
+from sklearn.metrics import (
+    accuracy_score, 
+    f1_score, 
+    recall_score,
+    balanced_accuracy_score
+)
+
+from glam4cm.tokenization.utils import get_tokenizer
+from glam4cm.utils import merge_argument_parsers, set_seed
+
+
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+    acc = (preds == labels).mean()
+    f1_macro = f1_score(labels, preds, average='macro')
+    accuracy = accuracy_score(labels, preds)
+    recall = recall_score(labels, preds, average='macro')
+    balanced_acc = balanced_accuracy_score(labels, preds)
+
+    return {
+        'balanced_accuracy': balanced_acc,
+        'accuracy': acc,
+        'f1_macro': f1_macro,
+        'precision': accuracy,
+        'recall': recall
+    }
+
+
+def get_num_labels(dataset):
+    train_labels = dataset['train'][:]['labels'].unique().tolist()
+    test_labels = dataset['test'][:]['labels'].unique().tolist()
+    return len(set(train_labels + test_labels))
+
+
+def get_parser():
+    common_parser = get_common_args_parser()
+    bert_parser = get_bert_args_parser()
+    parser = merge_argument_parsers(common_parser, bert_parser)
+
+    parser.add_argument('--oversampling_ratio', type=float, default=-1)
+
+    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
+    distance = args.distance
+    dataset = get_models_dataset(dataset_name, **config_params)
+
+    print("Loaded dataset")
+
+    graph_data_params = get_config_params(args)
+    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()
+
+
+if __name__ == '__main__':
+    args = get_parser()
+    run(args)

glam4cm/downstream_tasks/cm_gpt_edge_classification.py

@@ -0,0 +1,73 @@
+import os
+from glam4cm.downstream_tasks.common_args import (
+    get_common_args_parser, 
+    get_gpt_args_parser
+)
+
+from glam4cm.data_loading.graph_dataset import GraphEdgeDataset
+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, default='type')
+    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 = GraphEdgeDataset(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_link_prediction_lm_data(
+        tokenizer=tokenizer,
+        label=args.node_cls_label
+    )
+
+    if args.pretr and os.path.exists(args.pretr):
+        cmgpt = CMGPT.from_pretrained(f"{args.pretr}")
+    else:
+        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()