Readme update

config.json

@@ -1,5 +1,12 @@
 {
+  "architectures": [
+    "VCGNModelForTextClassification"
+  ],
   "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoConfig": "configuration_vgcn.VGCNConfig",
+    "AutoModelForSequenceClassification": "modeling_vcgn.VCGNModelForTextClassification"
+  },
   "bert_model": "readerbench/RoBERT-base",
   "classifier_dropout": null,
   "do_lower_case": 1,
@@ -34,6 +41,7 @@
   "pad_token_id": 0,
   "position_embedding_type": "absolute",
   "tf_threshold": 0.0,
+  "torch_dtype": "float32",
   "transformers_version": "4.31.0",
   "type_vocab_size": 2,
   "use_cache": true,

configuration_vgcn.py

@@ -6,6 +6,7 @@ class VGCNConfig(BertConfig):
 
     def __init__(
         self,
+        bert_model='readerbench/RoBERT-base',
         gcn_adj_matrix: str ='',
         max_seq_len: int = 256,
         npmi_threshold: float = 0.2,
@@ -29,5 +30,6 @@ class VGCNConfig(BertConfig):
         self.tf_threshold = tf_threshold
         self.vocab_type = vocab_type
         self.gcn_embedding_dim = gcn_embedding_dim
+        self.bert_model = bert_model
 
         super().__init__(**kwargs)

modeling_vcgn.py

@@ -64,27 +64,51 @@ def get_torch_gcn(gcn_vocab_adj_tf, gcn_vocab_adj,gcn_config:VGCNConfig):
 class VCGNModelForTextClassification(PreTrainedModel):
     config_class = VGCNConfig
 
-    def __init__(self, config):
+    def __init__(self, config, load_adjacency_matrix=True,):
         super().__init__(config)
-
-        self.pre_trained_model_name = ''
-        self.remove_stop_words = False
-        self.tokenizer = None
-        self.norm_gcn_vocab_adj_list = None
-        self.gcn_vocab_size = config.vocab_size
            
+        self.tokenizer = BertTokenizer.from_pretrained(config.bert_model)
 
-        self.load_adj_matrix(config.gcn_adj_matrix)
+        if load_adjacency_matrix:
+            norm_gcn_vocab_adj_list = self.load_adj_matrix(config.gcn_adj_matrix)
+        else:
+            norm_gcn_vocab_adj_list = []
+            for _ in range(2 if config.vocab_type=='all' else 1):
+                norm_gcn_vocab_adj_list.append(torch.sparse.FloatTensor(torch.LongTensor([[0],[0]]), torch.Tensor([0]), (config.vocab_size, config.vocab_size)))
 
         self.model = VGCN_Bert(
           config,
-          gcn_adj_matrix=self.norm_gcn_vocab_adj_list, 
+          gcn_adj_matrix=norm_gcn_vocab_adj_list, 
           gcn_adj_dim=config.vocab_size, 
-          gcn_adj_num=len(self.norm_gcn_vocab_adj_list), 
+          gcn_adj_num=len(norm_gcn_vocab_adj_list), 
           gcn_embedding_dim=config.gcn_embedding_dim,
 
         )
 
+    @classmethod
+    def from_pretrained(cls, *model_args, reload_adjacency_matrix=False, **kwargs):
+        model = super().from_pretrained( *model_args, **kwargs, load_adjacency_matrix=False)
+
+        if reload_adjacency_matrix:
+            norm_gcn_vocab_adj_list = model.load_adj_matrix(model.config.gcn_adj_matrix)
+            model.model.embeddings.vocab_gcn.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in norm_gcn_vocab_adj_list])
+            for p in model.model.embeddings.vocab_gcn.adj_matrix:
+                p.requires_grad=False
+
+        return model
+    
+    def set_adjacency_matrix(self, adj_matrix:Union[List, np.ndarray, sp.csr_matrix, torch.Tensor] ):
+        
+        if isinstance(adj_matrix, np.ndarray):
+            adj_matrix = [torch.from_numpy(adj_matrix)]
+        else:
+            raise ValueError(f"adjacency matrix must be a list of torch.Tensor or torch.nn.Parameter, got {type(adj_matrix)}")
+
+        self.model.embeddings.vocab_gcn.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in adj_matrix])
+        for p in self.model.embeddings.vocab_gcn.adj_matrix:
+            p.requires_grad=False
+
+
     def load_adj_matrix(self, adj_matrix):
         filename = None
         if Path(adj_matrix).is_file():
@@ -98,11 +122,8 @@ class VCGNModelForTextClassification(PreTrainedModel):
 
         gcn_vocab_adj_tf, gcn_vocab_adj, adj_config = pkl.load(open(filename, 'rb'))
 
-
-        self.pre_trained_model_name = adj_config['bert_model']
-        self.remove_stop_words = adj_config['remove_stop_words']
-        self.tokenizer =  BertTokenizer.from_pretrained(self.pre_trained_model_name)
-        self.norm_gcn_vocab_adj_list =  get_torch_gcn(gcn_vocab_adj_tf, gcn_vocab_adj, self.config)
+        self.tokenizer =  BertTokenizer.from_pretrained(adj_config['bert_model'])
+        return get_torch_gcn(gcn_vocab_adj_tf, gcn_vocab_adj, self.config)
 
     def _prep_batch(self, batch: torch.Tensor):
 
@@ -207,12 +228,14 @@ class VocabGraphConvolution(nn.Module):
     """
     def __init__(self,adj_matrix,voc_dim, num_adj, hid_dim, out_dim, dropout_rate=0.2):
         super(VocabGraphConvolution, self).__init__()
-        if type(adj_matrix) is not list:
+        if isinstance(adj_matrix, nn.Parameter) or isinstance(adj_matrix, nn.ParameterList):
             self.adj_matrix=adj_matrix
-        else:
+        elif isinstance(adj_matrix, list):
             self.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in adj_matrix])
             for p in self.adj_matrix:
                 p.requires_grad=False
+        else:
+            raise ValueError(f"adjacency matrix must be a list of torch.Tensor or torch.nn.Parameter, got {type(adj_matrix)}")
 
         self.voc_dim=voc_dim
         self.num_adj=num_adj