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

glam4cm/downstream_tasks/bert_edge_classification.py

@@ -2,10 +2,15 @@ import os
 from transformers import TrainingArguments, Trainer
 from glam4cm.data_loading.graph_dataset import GraphEdgeDataset
 from glam4cm.data_loading.utils import oversample_dataset
-from glam4cm.settings import LP_TASK_EDGE_CLS
-from glam4cm.downstream_tasks.common_args import get_bert_args_parser, get_common_args_parser, get_config_params
+from glam4cm.settings import EDGE_CLS_TASK, results_dir
+from glam4cm.downstream_tasks.common_args import (
+    get_bert_args_parser, 
+    get_common_args_parser, 
+    get_config_params,
+    get_config_str
+)
 from glam4cm.models.hf import get_model
-from glam4cm.downstream_tasks.utils import get_models_dataset
+from glam4cm.downstream_tasks.utils import get_logging_steps, get_models_dataset
 
 
 from sklearn.metrics import (
@@ -16,7 +21,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
 
 
 def compute_metrics(pred):
@@ -49,27 +54,39 @@ def get_parser():
 
 def run(args):
     set_seed(args.seed)
+    dataset_name = args.dataset
+    output_dir = os.path.join(
+        results_dir,
+        dataset_name,
+        f"LM_{EDGE_CLS_TASK}",
+        f'{args.edge_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,
         language = args.language,
         reload=args.reload
     )
-    dataset_name = args.dataset
+    
     
     print("Loaded 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}
+    graph_data_params = {**graph_data_params, 'task_type': EDGE_CLS_TASK}
 
     print("Loading graph dataset")
     graph_dataset = GraphEdgeDataset(dataset, **graph_data_params)
     print("Loaded graph dataset")
 
-    assert hasattr(graph_dataset, f'num_edges_{args.edge_cls_label}'), f"Dataset does not have node_{args.edge_cls_label} attribute"
+    assert hasattr(graph_dataset, f'num_edges_{args.edge_cls_label}'), f"Dataset does not have edge_{args.edge_cls_label} attribute"
     num_labels = getattr(graph_dataset, f"num_edges_{args.edge_cls_label}")
 
 
@@ -78,6 +95,11 @@ def run(args):
 
     print("Getting Edge Classification data")
     bert_dataset = graph_dataset.get_link_prediction_lm_data(tokenizer=tokenizer)
+    
+    train_dataset = bert_dataset['train']
+    test_dataset = bert_dataset['test']
+    set_encoded_labels(train_dataset, test_dataset)
+
 
     # exit(0)
     
@@ -88,28 +110,32 @@ def run(args):
     print("Training model")
     print(f'Number of labels: {num_labels}')
     
-    model = get_model(args.ckpt if args.ckpt else model_name, num_labels, len(tokenizer))
-    
+    model = get_model(
+        args.ckpt if args.ckpt else model_name, 
+        num_labels, 
+        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
 
-    output_dir = os.path.join(
-        'results',
-        dataset_name,
-        'edge_cls',
-        f'{args.edge_cls_label}',
-        f"{graph_dataset.config_hash}",
-    )
 
     logs_dir = os.path.join(
         'logs',
         dataset_name,
-        'edge_cls',
+        f"LM_{EDGE_CLS_TASK}",
         f'{args.edge_cls_label}',
         f"{graph_dataset.config_hash}",
     )
 
+    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,
@@ -117,13 +143,14 @@ def run(args):
         per_device_eval_batch_size=args.eval_batch_size,
         weight_decay=0.01,
         logging_dir=logs_dir,
-        logging_steps=args.num_log_steps,
+        logging_steps=logging_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,
+        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(

glam4cm/downstream_tasks/bert_graph_classification.py

@@ -12,10 +12,16 @@ from transformers import (
 
 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.downstream_tasks.common_args import (
+    get_bert_args_parser, 
+    get_common_args_parser, 
+    get_config_params,
+    get_config_str
+)
+from glam4cm.downstream_tasks.utils import get_logging_steps, get_models_dataset
+from glam4cm.settings import GRAPH_CLS_TASK, results_dir
 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
 
 
 
@@ -51,6 +57,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,
@@ -60,7 +67,7 @@ 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}
     print("Loading graph dataset")
     graph_dataset = GraphNodeDataset(dataset, **graph_data_params)
     print("Loaded graph dataset")
@@ -76,30 +83,50 @@ def run(args):
         train_dataset = classification_dataset['train']
         test_dataset = classification_dataset['test']
         num_labels = classification_dataset['num_classes']
+        
+        set_encoded_labels(train_dataset, test_dataset)
 
         print(len(train_dataset), len(test_dataset), num_labels)
 
         print("Training model")
         output_dir = os.path.join(
-            'results',
+            results_dir,
             dataset_name,
-            f'graph_cls_',
-            f"{graph_dataset.config_hash}",
+            f"LM_{GRAPH_CLS_TASK}",
+            f'{args.cls_label}',
+            get_config_str(args)
         )
+        # if os.path.exists(output_dir):
+        #     print(f"Output directory {output_dir} already exists. Exiting.")
+        #     exit(0)
 
         logs_dir = os.path.join(
             'logs',
             dataset_name,
-            f'graph_cls_',
-            f"{graph_dataset.config_hash}"
+            f"LM_{GRAPH_CLS_TASK}",
+            f'{args.cls_label}',
+            f"{graph_dataset.config_hash}_{fold_id}",
+            
         )
 
-        model = get_model(args.ckpt if args.ckpt else model_name, num_labels, len(tokenizer))
+        model = get_model(
+            args.ckpt if args.ckpt else model_name, 
+            num_labels, 
+            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
 
+
+        logging_steps = get_logging_steps(
+            len(train_dataset), 
+            args.num_epochs, 
+            args.train_batch_size
+        )
+#         
         # Training arguments
         training_args = TrainingArguments(
             output_dir=output_dir,
@@ -111,12 +138,13 @@ def run(args):
             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,
+            logging_steps=logging_steps,
+            eval_steps=logging_steps,
+            save_steps=logging_steps,
             save_total_limit=2,
             load_best_model_at_end=True,
-            fp16=True
+            fp16=True,
+            save_strategy="steps"
         )
 
         # Trainer
@@ -133,5 +161,7 @@ def run(args):
         results = trainer.evaluate()
         print(results)
         
+        trainer.save_model()
+        
         fold_id += 1
         break

glam4cm/downstream_tasks/bert_graph_classification_comp.py

@@ -1,3 +1,4 @@
+from collections import Counter
 import os
 import json
 from argparse import ArgumentParser
@@ -20,8 +21,7 @@ from transformers import (
 
 from glam4cm.data_loading.encoding import EncodingDataset
 from glam4cm.models.hf import get_model
-
-
+from glam4cm.settings import results_dir
 
 def compute_metrics(pred):
     labels = pred.label_ids
@@ -49,13 +49,17 @@ def get_parser():
     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('--limit', type=int, default=-1)
+    parser.add_argument('--trust_remote_code', action='store_true')
+    parser.add_argument('--include_dummies', action='store_true')
+    parser.add_argument('--task_type', type=str, default='graph_cls')
 
     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('--num_log_steps', type=int, default=500)
+    parser.add_argument('--num_eval_steps', type=int, default=500)
+    parser.add_argument('--num_save_steps', type=int, default=500)
     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)
@@ -66,15 +70,27 @@ def get_parser():
 def run(args):
     dataset_name = args.dataset_name
     model_name = args.model_name
-
-
+    include_dummies = args.include_dummies
+    
+    file_name = 'ecore.jsonl' if include_dummies and dataset_name == 'modelset' else 'ecore-with-dummy.jsonl'
     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')
+        if g['labels'] not in ['dummy', 'unknown']
+    ]
+    allowed_labels = [label for label, _ in dict(Counter([t[1] for t in texts]).most_common(48)).items()]
+    texts = [
+        (t, l) for t, l in texts
+        if l in allowed_labels
     ]
+    if args.task_type == 'dd':
+        print("Task type: DD")
+        texts = [(t, 0 if l in 'dummy' else 1) for t, l in texts]
+        
     shuffle(texts)
+    limit = args.limit if args.limit > 0 else len(texts)
+    texts = texts[:limit]
+        
     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]
@@ -84,13 +100,14 @@ def run(args):
 
     num_labels = len(y_map)
 
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=args.trust_remote_code)
     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)):
-
+        # if i == 0:
+        #     continue
         print(f'Fold {i+1}/{k}')
 
         train_texts = [texts[i] for i in train_idx]
@@ -101,27 +118,34 @@ def run(args):
 
         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)
+        train_dataset = EncodingDataset(tokenizer, train_texts, train_y, max_length=args.max_length)
+        test_dataset = EncodingDataset(tokenizer, test_texts, test_y, max_length=args.max_length)
+        # import code; code.interact(local=locals())
 
-        model = get_model(args.ckpt if args.ckpt else model_name, num_labels, len(tokenizer))
+        model = get_model(
+            args.ckpt if args.ckpt else model_name, 
+            num_labels, 
+            len(tokenizer), 
+            trust_remote_code=args.trust_remote_code
+        )
 
         print("Training model")
         output_dir = os.path.join(
-            'results',
+            results_dir,
             dataset_name,
-            'graph_cls_comp',
+            f'graph_cls_comp_{"dummy" if include_dummies else ""}{i+1}',
         )
 
         logs_dir = os.path.join(
             'logs',
-            dataset_name,
-            'graph_cls_comp',
+            f"{dataset_name}_{args.model_name if args.ckpt is None else args.ckpt.split('/')[-1]}",
+            f'graph_cls_comp_{"dummy" if include_dummies else ""}{i+1}',
         )
 
         print("Running epochs: ", args.num_epochs)
 
         # Training arguments
+        print("Batch size: ", args.train_batch_size)
         training_args = TrainingArguments(
             output_dir=output_dir,
             num_train_epochs=args.num_epochs,
@@ -131,11 +155,10 @@ def run(args):
             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
+            logging_steps=args.num_log_steps,
+            eval_steps=args.num_eval_steps,
+            fp16=True,
+            save_strategy="no"
         )
 
         # Trainer
@@ -153,4 +176,4 @@ def run(args):
         print(results)
 
         i += 1
-        
+        # break

glam4cm/downstream_tasks/bert_link_prediction.py

@@ -3,9 +3,14 @@ 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.settings import LINK_PRED_TASK, results_dir
+from glam4cm.downstream_tasks.common_args import (
+    get_bert_args_parser, 
+    get_common_args_parser, 
+    get_config_params,
+    get_config_str
+)
+from glam4cm.downstream_tasks.utils import get_logging_steps, get_models_dataset
 from glam4cm.tokenization.special_tokens import *
 
 
@@ -17,7 +22,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
 
 
 def compute_metrics(pred):
@@ -56,6 +61,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,
@@ -69,17 +75,19 @@ def run(args):
 
 
     graph_data_params = get_config_params(args)
-    graph_data_params = {**graph_data_params, 'task': LP_TASK_LINK_PRED}
+    graph_data_params = {
+        **graph_data_params, 
+        'add_negative_train_samples': True, 
+        'neg_sampling_ratio': args.neg_sampling_ratio,
+        'task_type': LINK_PRED_TASK
+    }
+    print(graph_data_params)
 
     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
-    ))
+        **graph_data_params
+    )
     print("Loaded graph dataset")
 
 
@@ -89,33 +97,44 @@ def run(args):
 
 
     print("Getting link prediction data")
-    bert_dataset = graph_dataset.get_link_prediction_lm_data(
-        tokenizer=tokenizer,
-        task_type=LP_TASK_LINK_PRED
-    )
+    bert_dataset = graph_dataset.get_link_prediction_lm_data(tokenizer=tokenizer)
+    train_dataset = bert_dataset['train']
+    test_dataset = bert_dataset['test']
+    set_encoded_labels(train_dataset, test_dataset)
+
 
     print("Training model")
-    model = get_model(args.ckpt if args.ckpt else model_name, num_labels=2, len_tokenizer=len(tokenizer))
+    model = get_model(args.ckpt if args.ckpt else model_name, num_labels=2, 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
 
-
+    
     output_dir = os.path.join(
-        'results',
+        results_dir,
         dataset_name,
-        'lp',
-        f"{graph_dataset.config_hash}",
+        f"LM_{LINK_PRED_TASK}",
+        get_config_str(args)
     )
+    if os.path.exists(output_dir):
+        print(f"Output directory {output_dir} already exists. Exiting.")
+        exit(0)
 
     logs_dir = os.path.join(
         'logs',
         dataset_name,
-        'lp',
+        f"LM_{LINK_PRED_TASK}",
         f"{graph_dataset.config_hash}",
     )
     
+    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,
@@ -123,13 +142,14 @@ def run(args):
         per_device_eval_batch_size=args.eval_batch_size,
         weight_decay=0.01,
         logging_dir=logs_dir,
-        logging_steps=200,
+        logging_steps=logging_steps,
         eval_strategy='steps',
-        eval_steps=200,
-        save_steps=200,
-        save_total_limit=2,
-        load_best_model_at_end=True,
+        eval_steps=logging_steps,
+        # save_steps=200,
+        # save_total_limit=2,
+        # load_best_model_at_end=True,
         fp16=True,
+        save_strategy="no"
     )
 
     trainer = Trainer(