Update demo_model.py

demo_model.py

@@ -10,6 +10,7 @@ import utils
 import networkx as nx
 from sentence_transformers import SentenceTransformer
 import pytorch_lightning as pl
+from transformers import BertTokenizer, BertForSequenceClassification
 
 
 class LGGMText2Graph_Demo(pl.LightningModule):
@@ -55,7 +56,7 @@ class LGGMText2Graph_Demo(pl.LightningModule):
         self.limit_dist = utils.PlaceHolder(X=x_limit, E=e_limit, y=y_limit)
         
 
-    def generate(self, text, num_nodes) -> None:
+    def generate_basic(self, text, num_nodes) -> None:
         print(num_nodes)
         prompt_emb = torch.tensor(self.text_encoder.encode([text])).to(self.device)
         samples = self.sample_batch(5, cond_emb = prompt_emb, num_nodes = num_nodes)
@@ -70,8 +71,33 @@ class LGGMText2Graph_Demo(pl.LightningModule):
         
         return nx_graphs
 
-    def init_prompt_encoder(self):
+    def generate_pretrained(self, text, num_nodes) -> None:
+        encoded_input = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)
+        encoded_input = {key: val.to(self.text_encoder.device) for key, val in encoded_input.items()}
+
+        # Get the model output
+        with torch.no_grad():
+            prompt_emb = self.text_encoder(**encoded_input).hidden_states[-1][:, 0]
+
+        samples = self.sample_batch(5, cond_emb = prompt_emb.to(self.device), num_nodes = num_nodes)
+
+        nx_graphs = []
+        for graph in samples:
+            node_types, edge_types = graph
+            A = edge_types.bool().cpu().numpy()
+
+            nx_graph = nx.from_numpy_array(A)
+            nx_graphs.append(nx_graph)
+        
+        return nx_graphs
+
+    def init_prompt_encoder_basic(self):
         self.text_encoder = SentenceTransformer("all-MiniLM-L6-v2")
+    
+    def init_prompt_encoder_pretrained(self):
+        model_name = f"./checkpoint-900"  # or "bert-base-uncased" if starting from the base model
+        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        self.text_encoder = BertForSequenceClassification.from_pretrained(model_name, num_labels=8, output_hidden_states=True, device_map = 'cpu')
 
 
     @torch.no_grad()