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

glam4cm/settings.py

@@ -9,6 +9,7 @@ logger.setLevel(logging.DEBUG)
 
 
 BERT_MODEL = 'bert-base-uncased'
+MODERN_BERT = 'answerdotai/ModernBERT-base'
 WORD2VEC_MODEL = 'word2vec'
 TFIDF_MODEL = 'tfidf'
 FAST_TEXT_MODEL = 'uml-fasttext.bin'
@@ -35,12 +36,22 @@ modelsets_ecore_json_path = os.path.join(datasets_dir, 'modelset/ecore.jsonl')
 
 
 graph_data_dir = 'datasets/graph_data'
+results_dir = 'results'
 
 # Path: settings.py
 
 
-LP_TASK_EDGE_CLS = 'edge_cls'
-LP_TASK_LINK_PRED = 'lp'
+EDGE_CLS_TASK = 'edge_cls'
+LINK_PRED_TASK = 'lp'
+NODE_CLS_TASK = 'node_cls'
+GRAPH_CLS_TASK = 'graph_cls'
+DUMMY_GRAPH_CLS_TASK = 'dummy_graph_cls'
+
+
+SEP = ' '
+REFERENCE = 'reference'
+SUPERTYPE = 'supertype'
+CONTAINMENT = 'containment'
 
 
 EPOCH = 'epoch'
@@ -52,3 +63,5 @@ TEST_ACC = 'test_acc'
 TRAINING_PHASE = 'train'
 VALIDATION_PHASE = 'val'
 TESTING_PHASE = 'test'
+
+

glam4cm/tokenization/special_tokens.py

@@ -1,4 +1,26 @@
 EDGE_START = '<edge_begin>'
 EDGE_END = '<edge_end>'
 NODE_BEGIN = '<node_begin>'    
-NODE_END = '<node_end>'
+NODE_END = '<node_end>'
+
+escape_keywords = [
+    "EString",
+    "EInt",
+    "EBoolean",
+    "EFloat",
+    "EAttribute",    
+    "EReference",
+    "EClass",
+    "EEnum",
+    "EEnumLiteral",
+    "EDataType",
+    "EOperation",
+    "EParameter",
+    "ETypeParameter",
+    "EAnnotation",
+    "stereotype",
+    EDGE_START,
+    EDGE_END,
+    NODE_BEGIN,
+    NODE_END
+]

glam4cm/tokenization/utils.py

@@ -1,6 +1,7 @@
 from re import finditer
+from typing import List
 from glam4cm.tokenization.special_tokens import (
-    EDGE_START, EDGE_END, NODE_BEGIN, NODE_END
+    EDGE_START, EDGE_END, NODE_BEGIN, NODE_END, escape_keywords
 )
 from transformers import AutoTokenizer
 
@@ -24,6 +25,8 @@ def get_tokenizer(model_name, use_special_tokens=False, max_length=512) -> AutoT
 
 
 def camel_case_split(identifier) -> list:
+    if any(ek in identifier for ek in escape_keywords):
+        return [identifier]
     matches = finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)', identifier)
     return [m.group(0) for m in matches]
 
@@ -31,7 +34,23 @@ def camel_case_split(identifier) -> list:
 def doc_tokenizer(doc, lower=False) -> str:
     words = doc.split()
     # split _
-    words = [w2 for w1 in words for w2 in w1.split('_') if w2 != '']
+    snake_words: List[str] = list()
+    for w1 in words:
+        if any(ek in w1 for ek in escape_keywords):
+            snake_words.append(w1)
+        else:
+            snake_words.extend([w2 for w2 in w1.split('_') if w2 != ''])
+            
+    
     # camelcase
-    words = [w2.lower() if lower else w2 for w1 in words for w2 in camel_case_split(w1) if w2 != '']
-    return " ".join(words)
+    final_words: List[str] = list()
+    for word in snake_words:
+        if any(ek in word for ek in escape_keywords):
+            final_words.append(word)
+        else:
+            final_words.extend(camel_case_split(word))
+            
+    if lower:
+        final_words = [w.lower() for w in final_words]
+        
+    return " ".join(final_words)

glam4cm/trainers/cm_gpt_trainer.py

@@ -52,7 +52,7 @@ class CMGPTTrainer:
             self.compute_metrics = compute_metrics
         
         print(f"Number of parameters: {sum(p.numel() for p in self.model.parameters() if p.requires_grad)/ 1000000:.3f}M")
-
+        print(f"Logging to: {log_dir}")
 
     def step(self, batch, idx=None):
         # B, T = batch['input_ids'].shape

glam4cm/trainers/gnn_edge_classifier.py

@@ -57,6 +57,9 @@ class GNNEdgeClassificationTrainer(Trainer):
         all_preds, all_labels = list(), list()
         epoch_loss = 0
         epoch_metrics = defaultdict(float)
+        # print("Total dataloader size: ", len(self.dataloader))
+        # from tqdm.auto import tqdm
+        # for data in tqdm(self.dataloader):
         for data in self.dataloader:
             self.optimizer.zero_grad()
             self.model.zero_grad()
@@ -81,7 +84,10 @@ class GNNEdgeClassificationTrainer(Trainer):
         
         all_preds = torch.cat(all_preds, dim=0)
         all_labels = torch.cat(all_labels, dim=0)
+        # import time
+        # t1 = time.time()
         epoch_metrics = self.compute_metrics(all_preds, all_labels)
+        # print(f"Time taken: {time.time() - t1}")
         epoch_metrics['loss'] = epoch_loss        
         epoch_metrics['phase'] = 'train'
 
@@ -99,11 +105,16 @@ class GNNEdgeClassificationTrainer(Trainer):
             epoch_metrics = defaultdict(float)
             for data in self.dataloader:
                 x = data.x
+                train_edge_index =  data.train_pos_edge_label_index
+                train_mask = data.train_edge_mask
+                train_edge_attr = data.edge_attr[train_mask] if self.use_edge_attrs else None
+                
                 edge_index =  data.test_pos_edge_label_index
                 test_mask = data.test_edge_mask
                 edge_attr = data.edge_attr[test_mask] if self.use_edge_attrs else None
                 
-                h = self.get_logits(x, edge_index, edge_attr)
+                h = self.get_logits(x, train_edge_index, train_edge_attr)
+                
                 scores = self.get_prediction_score(h, edge_index, edge_attr)
                 labels = getattr(data, f"edge_{self.cls_label}")[test_mask]
                 all_preds.append(scores.detach().cpu())

glam4cm/trainers/gnn_graph_classifier.py

@@ -1,4 +1,4 @@
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 import torch
 from collections import defaultdict
 from torch_geometric.loader import DataLoader
@@ -22,7 +22,7 @@ class GNNGraphClassificationTrainer(Trainer):
             self, 
             model: GNNConv, 
             predictor: GraphClassifer,
-            dataset: List[Tuple[Data, Data]],
+            dataset: Dict[str, List[Data]],
             cls_label='label',
             lr=1e-4,
             num_epochs=100,
@@ -43,8 +43,14 @@ class GNNGraphClassificationTrainer(Trainer):
 
         self.cls_label = cls_label
         self.dataloaders = dict()
-        self.dataloaders['train'] = DataLoader(dataset['train'], batch_size=batch_size, shuffle=True)
-        self.dataloaders['test'] = DataLoader(dataset['test'], batch_size=batch_size, shuffle=False)
+        self.dataloaders['train'] = DataLoader(
+            [g for g in dataset['train'] if len(g.edge_index) != 0], 
+            batch_size=batch_size, shuffle=True
+        )
+        self.dataloaders['test'] = DataLoader(
+            [g for g in dataset['test'] if len(g.edge_index) != 0], 
+            batch_size=batch_size, shuffle=False
+        )
 
         self.results = list()
 
@@ -120,4 +126,5 @@ class GNNGraphClassificationTrainer(Trainer):
         s2t = lambda x: x.replace("_", " ").title()
         print(f"Epoch: {len(self.results)//2} {' | '.join([f'{s2t(k)}: {v:.4f}' for k, v in epoch_metrics.items() if k != 'phase'])}")
 
-        return epoch_metrics
+        return epoch_metrics
+    

glam4cm/trainers/gnn_link_predictor.py

@@ -58,6 +58,12 @@ class GNNLinkPredictionTrainer(Trainer):
         all_preds, all_labels = list(), list()
         epoch_loss = 0
         epoch_metrics = defaultdict(float)
+        
+        total_pos_edges = sum([data.train_pos_edge_label_index.size(1) for data in self.dataloader.dataset])
+        total_neg_edges = sum([data.train_neg_edge_label_index.size(1) for data in self.dataloader.dataset])
+        print(f"Total positive edges: {total_pos_edges}")
+        print(f"Total negative edges: {total_neg_edges}")
+        
         for data in tqdm(self.dataloader, desc='Training Batches'):
             self.optimizer.zero_grad()
             self.model.zero_grad()
@@ -103,14 +109,23 @@ class GNNLinkPredictionTrainer(Trainer):
             for data in tqdm(self.dataloader, desc='Testing Batches'):
                 
                 x = data.x
+                
+                train_edge_index = torch.cat([
+                    data.train_pos_edge_label_index,
+                    data.train_neg_edge_label_index
+                ], dim=1)
+                train_edge_attr = (
+                    data.edge_attr[data.train_edge_mask]
+                    if self.use_edge_attrs else None
+                )
+                
+                h = self.get_logits(x, train_edge_index, train_edge_attr)
+                
                 pos_edge_index =  data.test_pos_edge_label_index
                 neg_edge_index = data.test_neg_edge_label_index
                 test_mask = data.test_edge_mask
                 edge_attr = data.edge_attr[test_mask] if self.use_edge_attrs else None
 
-
-                h = self.get_logits(x, pos_edge_index, edge_attr)
-                # h = x
                 pos_score = self.get_prediction_score(h, pos_edge_index, edge_attr)
                 neg_score = self.get_prediction_score(h, neg_edge_index, edge_attr)
 

glam4cm/trainers/gnn_link_predictor_v2.py

@@ -0,0 +1,146 @@
+from sklearn.metrics import roc_auc_score, average_precision_score
+from torch_geometric.loader import DataLoader
+from torch_geometric.data import Data
+from torch_geometric.nn import GATConv, VGAE
+import torch.nn.functional as F
+import torch
+from typing import List
+
+from glam4cm.models.gnn_layers import (
+    GNNConv, 
+    EdgeClassifer
+)
+
+from tqdm.auto import tqdm
+from glam4cm.settings import device
+
+
+class GATVGAEEncoder(torch.nn.Module):
+    def __init__(self, in_channels, hid_channels, out_channels, heads=(4,2), dropout=0.3):
+        super().__init__()
+        self.conv1 = GATConv(in_channels, hid_channels, heads=heads[0], dropout=dropout)
+        # mu and log_std each map to latent dim
+        self.conv_mu     = GATConv(hid_channels * heads[0], out_channels, heads=heads[1], concat=False)
+        self.conv_logstd = GATConv(hid_channels * heads[0], out_channels, heads=heads[1], concat=False)
+
+    def forward(self, x, edge_index):
+        x = F.elu(self.conv1(x, edge_index))
+        return self.conv_mu(x, edge_index), self.conv_logstd(x, edge_index)
+
+
+
+class GNNLinkPredictionTrainerV2:
+    """
+    Trainer class for GNN Link Prediction
+    This class is used to train the GNN model for the link prediction task
+    The model is trained to predict the link between two nodes
+    """
+    def __init__(
+            self, 
+            model: GNNConv, 
+            predictor: EdgeClassifer, 
+            dataset: List[Data],
+            cls_label='type',
+            lr=1e-3,
+            num_epochs=100,
+            batch_size=32,
+            use_edge_attrs=False,
+            logs_dir='./logs'
+        ) -> None:
+        
+        self.num_epochs = num_epochs
+        self.lr = lr
+        in_dim = dataset[0].data.x.shape[1]
+        hid_dim = 64
+        out_dim = 32
+        self.encoder = GATVGAEEncoder(in_dim, hid_dim, out_dim).to(device)
+        self.model   = VGAE(self.encoder).to(device)
+        self.opt     = torch.optim.Adam(model.parameters(), lr=lr)
+
+        self.dataloader = DataLoader(
+            dataset, batch_size=batch_size, shuffle=True
+        )
+        self.results = list()
+
+        print("GNN Trainer initialized.")
+
+
+
+    def train(self):
+        self.model.train()
+        total_loss = 0.0
+
+        for data in self.dataloader:
+            data = data.to(device)
+            # Encode over the **train positive** graph only:
+            z = self.model.encode(data.x, data.train_pos_edge_label_index)
+
+            # recon_loss only on positives:
+            loss = self.model.recon_loss(z, data.train_pos_edge_label_index)
+            # KL regularizer:
+            loss += (1. / data.num_nodes) * self.model.kl_loss()
+
+            self.opt.zero_grad()
+            loss.backward()
+            self.opt.step()
+
+            total_loss += loss.item()
+
+        return total_loss / len(self.dataloader)
+
+
+    @torch.no_grad()
+    def test(self):
+        self.model.eval()
+        all_auc, all_ap = [], []
+
+        for data in self.dataloader:
+            data = data.to(device)
+            z = self.model.encode(data.x, data.train_pos_edge_label_index)
+
+            # positive edges from your test split
+            pos_idx = data.test_pos_edge_label_index
+            # generate equal‐size negative sample
+            neg_idx = data.test_neg_edge_label_index
+
+            pos_scores = self.model.decoder(z, pos_idx).sigmoid()
+            neg_scores = self.model.decoder(z, neg_idx).sigmoid()
+
+            y_true  = torch.cat([torch.ones(pos_scores.size(0)),
+                                torch.zeros(neg_scores.size(0))]).cpu()
+            y_score = torch.cat([pos_scores, neg_scores]).cpu()
+
+            all_auc.append( roc_auc_score(y_true, y_score) )
+            all_ap.append(  average_precision_score(y_true, y_score) )
+
+        return {
+            'AUC': sum(all_auc) / len(all_auc),
+            'AP':  sum(all_ap)  / len(all_ap),
+        }
+        
+
+    def compute_loss(self, pos_score, neg_score):
+        pos_label = torch.ones(pos_score.size(0), dtype=torch.long).to(device)
+        neg_label = torch.zeros(neg_score.size(0), dtype=torch.long).to(device)
+
+        scores = torch.cat([pos_score, neg_score], dim=0)
+        labels = torch.cat([pos_label, neg_label], dim=0)
+
+        loss = self.criterion(scores, labels)
+        return loss
+    
+    
+    def run(self):
+        all_metrics = list()
+        for epoch in tqdm(range(self.num_epochs), desc="Running Epochs"):
+            self.train()
+            test_metrics = self.test()
+            all_metrics.append(test_metrics)
+            print(f"Epoch {epoch+1}/{self.num_epochs} | AUC: {test_metrics['AUC']:.4f} | AP: {test_metrics['AP']:.4f}")
+
+        print("Training complete.")
+        best_metrics = sorted(all_metrics, key=lambda x: x['AUC'], reverse=True)[0]
+        
+        s2t = lambda x: x.replace("_", " ").title()
+        print(f"Best: {' | '.join([f'{s2t(k)}: {v:.4f}' for k, v in best_metrics.items()])}")
+    

glam4cm/trainers/gnn_trainer.py

@@ -111,9 +111,11 @@ class Trainer:
 
 
     def run(self):
+        all_metrics = list()
         for epoch in tqdm(range(self.num_epochs), desc="Running Epochs"):
             train_metrics = self.train()
             test_metrics = self.test()
+            all_metrics.append(test_metrics)
 
             for k, v in train_metrics.items():
                 if k != 'phase':
@@ -124,6 +126,12 @@ class Trainer:
                     self.writer.add_scalar(f"test/{k}", v, epoch)
     
         self.writer.close()
+        print("Training complete.")
+        best_metrics = sorted(all_metrics, key=lambda x: x['balanced_accuracy'], reverse=True)[0]
+        
+        s2t = lambda x: x.replace("_", " ").title()
+        print(f"Best: {' | '.join([f'{s2t(k)}: {v:.4f}' for k, v in best_metrics.items() if k != 'phase'])}")
+        
     
     def compute_metrics(self, all_preds, all_labels):
         return compute_classification_metrics(all_preds, all_labels)

glam4cm/trainers/metrics.py

@@ -20,7 +20,7 @@ def compute_metrics(p):
     }
 
 
-def compute_classification_metrics(preds, labels):
+def compute_classification_metrics(preds: torch.Tensor, labels: torch.Tensor) -> dict:
     """
     Compute F1-score, balanced accuracy, precision, and recall for multi-class classification.