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

glam4cm/__init__.py

@@ -6,4 +6,5 @@ warnings.filterwarnings(
     category=UserWarning,
     module="pydantic.main",
 )
-__version__ = "0.1.0"
+__version__ = "0.1.0"
+

glam4cm/data_loading/data.py

@@ -19,12 +19,12 @@ from glam4cm.lang2graph.common import (
 
 from scipy.sparse import csr_matrix
 
-from glam4cm.settings import LP_TASK_EDGE_CLS, LP_TASK_LINK_PRED
+from glam4cm.settings import DUMMY_GRAPH_CLS_TASK, EDGE_CLS_TASK, LINK_PRED_TASK
 from glam4cm.tokenization.special_tokens import *
 from torch_geometric.transforms import RandomLinkSplit
 import torch
-from torch_geometric.data import Data, Dataset
-from typing import List, Optional, Sequence, Union
+from torch_geometric.data import Data
+from typing import List, Union
 from glam4cm.tokenization.utils import doc_tokenizer
 
 
@@ -98,6 +98,7 @@ class TorchGraph:
             self, 
             graph: Union[EcoreNxG, ArchiMateNxG], 
             metadata: Union[EcoreMetaData, ArchimateMetaData],
+            task_type: str,
             distance = 0,
             test_ratio=0.2,
             use_edge_types=False,
@@ -108,9 +109,12 @@ class TorchGraph:
             no_labels=False,
             node_cls_label=None,
             edge_cls_label='type',
+            
+            node_topk: List[Union[str, int]]=None,
             fp='test_graph.pkl'
         ):
 
+        self.task_type = task_type
         self.fp = fp
         self.graph = graph
         self.metadata = metadata
@@ -126,6 +130,8 @@ class TorchGraph:
         self.node_cls_label = node_cls_label
         self.edge_cls_label = edge_cls_label
         
+        self.node_topk = node_topk
+        
         self.distance = distance
         self.test_ratio = test_ratio
         self.data = NumpyData()
@@ -148,6 +154,10 @@ class TorchGraph:
         )
 
         edge_texts = self.get_graph_edge_strs()
+        
+        # print(f"Number of edges: {len(edge_texts)}")
+        # print("Edge strings: ", edge_texts[:50])
+        
         return node_texts, edge_texts
     
 
@@ -162,17 +172,20 @@ class TorchGraph:
 
         def generate_embeddings():
             if randomize_ne or embedder is None:
-                print("Randomizing node embeddings")
+                # print("Randomizing node embeddings")
                 self.data.x = np.random.randn(self.graph.number_of_nodes(), random_embed_dim)
             else:
                 self.data.x = embedder.embed(list(self.node_texts.values()))
             
             if randomize_ee or embedder is None:
-                print("Randomizing edge embeddings")
+                # print("Randomizing edge embeddings")
                 self.data.edge_attr = np.random.randn(self.graph.number_of_edges(), random_embed_dim)
             else:
-                self.data.edge_attr = embedder.embed(list(self.edge_texts.values()))
+                edge_texts = list(self.edge_texts.values())
+                self.data.edge_attr = embedder.embed(edge_texts) \
+                    if len(edge_texts) > 0 else np.empty((self.graph.number_of_edges(), random_embed_dim))
 
+        
         if os.path.exists(f"{self.fp}") and not reload:
             with open(f"{self.fp}", 'rb') as f:
                 obj: Union[TorchEdgeGraph, TorchNodeGraph] = pickle.load(f)
@@ -199,6 +212,7 @@ class TorchGraph:
             distance = self.distance
 
         subgraph = create_graph_from_edge_index(self.graph, edge_index)
+
         return get_node_texts(
             subgraph, 
             distance,
@@ -230,6 +244,7 @@ class TorchGraph:
                 self.graph.numbered_graph, 
                 (u, v), 
                 d=self.distance, 
+                task_type=self.task_type,
                 metadata=self.metadata, 
                 use_node_attributes=self.use_attributes, 
                 use_node_types=self.use_node_types, 
@@ -249,14 +264,13 @@ class TorchGraph:
     def validate_data(self):
         assert self.data.num_nodes == self.graph.number_of_nodes()
     
+    
     def set_graph_label(self):
         if self.metadata.graph_label is not None and not hasattr(self.graph, self.metadata.graph_label):  #Graph has a label
             text = doc_tokenizer("\n".join(list(self.node_texts.values())))
-            # print("Text:", text)
-            # print("-" * 100)
             setattr(self.graph, self.metadata.graph_label, text)
         
-
+        
     @property
     def name(self):
         return '.'.join(self.graph.graph_id.replace('/', '_').split('.')[:-1])
@@ -268,24 +282,28 @@ class TorchEdgeGraph(TorchGraph):
             self, 
             graph: Union[EcoreNxG, ArchiMateNxG], 
             metadata: Union[EcoreMetaData, ArchimateMetaData],
-            distance = 1,
-            test_ratio=0.2,
-            add_negative_train_samples=False,
-            neg_samples_ratio=1,
-            use_edge_types=False,
-            use_node_types=False,
-            use_edge_label=False,
-            use_attributes=False,
-            use_special_tokens=False,
-            node_cls_label=None,
-            edge_cls_label='type',
-            no_labels=False,
+            task_type: str,
+            distance: int  = 1,
+            test_ratio: float =0.2,
+            add_negative_train_samples: bool =False,
+            neg_samples_ratio: int =1,
+            use_edge_types: bool =False,
+            use_node_types: bool =False,
+            use_edge_label: bool =False,
+            use_attributes: bool =False,
+            use_special_tokens: bool =False,
+            node_cls_label: str =None,
+            edge_cls_label: str ='type',
+            no_labels: bool =False,
+            
+            node_topk: List[Union[str, int]]=None,
             fp: str = 'test_graph.pkl'
         ):
 
         super().__init__(
             graph=graph, 
             metadata=metadata, 
+            task_type=task_type,
             distance=distance, 
             test_ratio=test_ratio, 
             use_node_types=use_node_types,
@@ -296,6 +314,7 @@ class TorchEdgeGraph(TorchGraph):
             no_labels=no_labels,
             node_cls_label=node_cls_label,
             edge_cls_label=edge_cls_label,
+            node_topk=node_topk,
             fp=fp
         )
         self.add_negative_train_samples = add_negative_train_samples
@@ -304,9 +323,7 @@ class TorchEdgeGraph(TorchGraph):
         self.validate_data()
         self.set_graph_label()
 
-        
 
-    
     def get_pyg_data(self):
         
         d = GraphData()
@@ -344,7 +361,8 @@ class TorchEdgeGraph(TorchGraph):
         setattr(d, 'test_pos_edge_label', test_data.pos_edge_label)
 
 
-        if hasattr(train_data, 'neg_edge_label_index'):
+        if self.add_negative_train_samples:
+            assert hasattr(train_data, 'neg_edge_label_index')
             assert not any([self.graph.numbered_graph.has_edge(*edge) for edge in train_data.neg_edge_label_index.t().tolist()])
             assert not any([self.graph.numbered_graph.has_edge(*edge) for edge in test_data.neg_edge_label_index.t().tolist()])
             setattr(d, 'train_neg_edge_label_index', train_data.neg_edge_label_index)
@@ -372,6 +390,9 @@ class TorchEdgeGraph(TorchGraph):
         node_texts, edge_texts = self.get_node_edge_strings(
             edge_index=edge_index.numpy(),
         )
+        
+        # print("Node texts: ", list(node_texts.values())[:5])
+        # print("Edge texts: ", list(edge_texts.values())[:5])
 
         setattr(d, 'num_nodes', self.graph.number_of_nodes())
         setattr(d, 'num_edges', self.graph.number_of_edges())
@@ -384,7 +405,7 @@ class TorchEdgeGraph(TorchGraph):
         train_pos_edge_index = self.data.edge_index
         test_pos_edge_index = self.data.test_pos_edge_label_index
 
-        if task_type == LP_TASK_LINK_PRED:
+        if task_type == LINK_PRED_TASK:
             train_neg_edge_index = self.data.train_neg_edge_label_index
             test_neg_edge_index = self.data.test_neg_edge_label_index
         else:
@@ -412,9 +433,12 @@ class TorchEdgeGraph(TorchGraph):
             
             edge_strs = list(edge_strs.values())
             data[f'{edge_index_label}_edges'] = edge_strs
+            
+            # print(f"Number of {edge_index_label} edges: {len(edge_strs)}")
+            # print("Edge strings: ", edge_strs[:50])
 
 
-        if task_type == LP_TASK_EDGE_CLS and not only_texts:
+        if task_type == EDGE_CLS_TASK and not only_texts:
             train_mask = self.data.train_edge_mask
             test_mask = self.data.test_edge_mask
             train_classes, test_classes = getattr(self.data, f'edge_{label}')[train_mask], getattr(self.data, f'edge_{label}')[test_mask]
@@ -431,22 +455,28 @@ class TorchNodeGraph(TorchGraph):
             self, 
             graph: Union[EcoreNxG, ArchiMateNxG], 
             metadata: dict,
-            distance = 1,
-            test_ratio=0.2,
-            use_node_types=False,
-            use_edge_types=False,
-            use_edge_label=False,
-            use_attributes=False,
-            use_special_tokens=False,
-            no_labels=False,
-            node_cls_label=None,
-            edge_cls_label='type',
-            fp='test_graph.pkl'
+            task_type: str,
+            
+            distance: int = 1,
+            test_ratio: float =0.2,
+            use_node_types: bool =False,
+            use_edge_types: bool =False,
+            use_edge_label: bool =False,
+            use_attributes: bool =False,
+            use_special_tokens: bool =False,
+            no_labels: bool =False,
+            node_cls_label: str =None,
+            edge_cls_label: str ='type',
+            
+            node_topk: List[Union[str, int]]=None,
+            
+            fp='test_graph.pkl',
         ):
 
         super().__init__(
             graph, 
             metadata=metadata, 
+            task_type=task_type,
             distance=distance, 
             test_ratio=test_ratio, 
             use_node_types=use_node_types,
@@ -457,6 +487,8 @@ class TorchNodeGraph(TorchGraph):
             no_labels=no_labels,
             node_cls_label=node_cls_label,
             edge_cls_label=edge_cls_label,
+            
+            node_topk=node_topk,
             fp=fp
         )
         
@@ -468,15 +500,28 @@ class TorchNodeGraph(TorchGraph):
     
     def get_pyg_data(self):
         d = GraphData()
-        train_nodes, test_nodes = train_test_split(
-            list(self.graph.numbered_graph.nodes), 
-            test_size=self.test_ratio, 
-            shuffle=True, 
-            random_state=42
-        )
+        if self.task_type == DUMMY_GRAPH_CLS_TASK:
+            train_nodes = list(self.graph.numbered_graph.nodes)
+            test_nodes = list()
+        else:
+            train_nodes, test_nodes = train_test_split(
+                list(self.graph.numbered_graph.nodes), 
+                test_size=self.test_ratio, 
+                shuffle=True, 
+                random_state=42
+            )
 
+        def get_node_label(node):
+            if self.node_cls_label in self.graph.numbered_graph.nodes[node]\
+                and self.graph.numbered_graph.nodes[node][self.node_cls_label] is not None:
+                return self.graph.numbered_graph.nodes[node][self.node_cls_label]
+            return None
+        
         nx.set_node_attributes(self.graph.numbered_graph, {node: False for node in train_nodes}, 'masked')
-        nx.set_node_attributes(self.graph.numbered_graph, {node: True for node in test_nodes}, 'masked')
+        nx.set_node_attributes(self.graph.numbered_graph, {
+            node: get_node_label(node) in self.node_topk 
+            for node in test_nodes
+        }, 'masked')
 
         train_idx = torch.tensor(train_nodes, dtype=torch.long)
         test_idx = torch.tensor(test_nodes, dtype=torch.long)
@@ -527,105 +572,4 @@ def validate_edges(graph: Union[TorchEdgeGraph, TorchNodeGraph]):
     
     if train_neg_edge_index is not None and test_neg_edge_index is not None:
         assert len(set((a, b) for a, b in train_neg_edge_index.T.tolist()).intersection(set((a, b) for a, b in test_neg_edge_index.T.tolist()))) == 0
-    
-    
-
-
-class LinkPredictionCollater:
-    def __init__(
-            self, 
-            follow_batch: Optional[List[str]] = None, 
-            exclude_keys: Optional[List[str]] = None
-        ):
-        self.follow_batch = follow_batch
-        self.exclude_keys = exclude_keys
-
-    def __call__(self, batch: List[Data]):
-        # Initialize lists to collect batched properties
-        x = []
-        edge_index = []
-        edge_attr = []
-        y = []
-        overall_edge_index = []
-        edge_classes = []
-        train_edge_mask = []
-        test_edge_mask = []
-        train_pos_edge_label_index = []
-        train_pos_edge_label = []
-        train_neg_edge_label_index = []
-        train_neg_edge_label = []
-        test_pos_edge_label_index = []
-        test_pos_edge_label = []
-        test_neg_edge_label_index = []
-        test_neg_edge_label = []
-        
-        # Offsets for edge indices
-        node_offset = 0
-        edge_offset = 0
-
-        for data in batch:
-            x.append(data.x)
-            edge_index.append(data.edge_index + node_offset)
-            edge_attr.append(data.edge_attr)
-            y.append(data.y)
-            overall_edge_index.append(data.overall_edge_index + edge_offset)
-            edge_classes.append(data.edge_classes)
-
-            train_edge_mask.append(data.train_edge_mask)
-            test_edge_mask.append(data.test_edge_mask)
-
-            train_pos_edge_label_index.append(data.train_pos_edge_label_index + node_offset)
-            train_pos_edge_label.append(data.train_pos_edge_label)
-            train_neg_edge_label_index.append(data.train_neg_edge_label_index + node_offset)
-            train_neg_edge_label.append(data.train_neg_edge_label)
-
-            test_pos_edge_label_index.append(data.test_pos_edge_label_index + node_offset)
-            test_pos_edge_label.append(data.test_pos_edge_label)
-            test_neg_edge_label_index.append(data.test_neg_edge_label_index + node_offset)
-            test_neg_edge_label.append(data.test_neg_edge_label)
-
-            node_offset += data.num_nodes
-            edge_offset += data.edge_attr.size(0)
-
-        return GraphData(
-            x=torch.cat(x, dim=0),
-            edge_index=torch.cat(edge_index, dim=1),
-            edge_attr=torch.cat(edge_attr, dim=0),
-            y=torch.tensor(y),
-            overall_edge_index=torch.cat(overall_edge_index, dim=1),
-            edge_classes=torch.cat(edge_classes),
-            train_edge_mask=torch.cat(train_edge_mask),
-            test_edge_mask=torch.cat(test_edge_mask),
-            train_pos_edge_label_index=torch.cat(train_pos_edge_label_index, dim=1),
-            train_pos_edge_label=torch.cat(train_pos_edge_label),
-            train_neg_edge_label_index=torch.cat(train_neg_edge_label_index, dim=1),
-            train_neg_edge_label=torch.cat(train_neg_edge_label),
-            test_pos_edge_label_index=torch.cat(test_pos_edge_label_index, dim=1),
-            test_pos_edge_label=torch.cat(test_pos_edge_label),
-            test_neg_edge_label_index=torch.cat(test_neg_edge_label_index, dim=1),
-            test_neg_edge_label=torch.cat(test_neg_edge_label),
-            num_nodes=node_offset
-        )
-
-
-class LinkPredictionDataLoader(torch.utils.data.DataLoader):
-    def __init__(
-        self,
-        dataset: Union[Dataset, Sequence[Data]],
-        batch_size: int = 1,
-        shuffle: bool = False,
-        collate_fn=None,
-        follow_batch: Optional[List[str]] = None,
-        exclude_keys: Optional[List[str]] = None,
-        **kwargs,
-    ):
-        if collate_fn is None:
-            collate_fn = LinkPredictionCollater(follow_batch, exclude_keys)
-        
-        super().__init__(
-            dataset,
-            batch_size,
-            shuffle,
-            collate_fn=collate_fn,
-            **kwargs,
-        )
+    

glam4cm/data_loading/encoding.py

@@ -1,24 +1,34 @@
+from typing import List, Union
 from torch.utils.data import Dataset
 import torch
+from transformers import AutoTokenizer
+
+def get_max_length(tokenizer):
+    tokenizer_name = tokenizer.name_or_path.lower()
+    if 'modernbert' in tokenizer_name:
+        return 8000
+    return 512
 
 # Create your dataset
 class EncodingDataset(Dataset):
     def __init__(
             self, 
-            tokenizer, 
-            texts, 
-            labels=None, 
+            tokenizer: AutoTokenizer, 
+            texts: List[str], 
+            labels:List[Union[str, int]]=None, 
             max_length=512,
             remove_duplicates=False
         ):
+        
+        max_length = get_max_length(tokenizer)
 
         if remove_duplicates:
-            # print(f'Dataset with {len(texts)} samples before removing duplicates')
+            print(f'Dataset with {len(texts)} samples before removing duplicates')
             texts_to_id = {text: i for i, text in enumerate(texts)}
             texts = list(texts_to_id.keys())
             labels = [labels[i] for i in texts_to_id.values()] if labels else None            
         
-        # print(f'Created dataset with {len(texts)} samples')
+        # print(f'Encoding started with {len(texts)} samples')
 
         self.inputs = tokenizer(
             texts, 
@@ -31,7 +41,8 @@ class EncodingDataset(Dataset):
         if labels is not None:
             self.inputs['labels'] = torch.tensor(labels, dtype=torch.long) if labels is not None else None
 
-        print("Encoding Dataset created with {} samples".format(len(self.inputs['input_ids'])))
+        # print("Embedding shape: ", self.inputs['input_ids'].shape)
+        # print("Encoding Dataset created with {} samples".format(len(self.inputs['input_ids'])))
         # print("\n".join([f"Label: {l}, Text: {i}" for i, l in zip(texts, labels)]))
         # import code; code.interact(local=locals())