working tag embeddings

embeddings/NextTagEmbedding.py

@@ -110,9 +110,6 @@ class NextTagEmbeddingTrainer:
                 total_loss += loss.item()
             losses.append(total_loss)
 
-    def get_tag_embedding(self, tag: str):
-        return self.model.embedding.weight[self.tag_to_ix[tag]]
-
     def to_tensorboard(self, run_name: str):
         """
         Write embedding to Tensorboard projector
@@ -130,7 +127,7 @@ class NextTagEmbeddingTrainer:
         model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
 
         # unpickle the tag_to_ix
-        with open('tag_to_ix_' + model_path, 'rb') as f:
+        with open(model_path.replace('tag-emb', f'tag_to_ix_tag-emb'), 'rb') as f:
             model.tag_to_ix = pickle.load(f)
 
         return model
@@ -157,6 +154,11 @@ class NextTagEmbedding(nn.Module):
         log_probs = F.log_softmax(out, dim=1)
         return log_probs
 
+    def get_tag_embedding(self, tag: str):
+        assert tag in self.tag_to_ix, "Tag not in vocabulary!"
+        assert self.tag_to_ix is not None, "Tag to index mapping not set!"
+        return self.embedding.weight[self.tag_to_ix[tag]]
+
 
 if __name__ == '__main__':
     tet = NextTagEmbeddingTrainer(context_length=2, emb_size=30, excluded_tags=['python'], database_path="../stackoverflow.db")

embeddings/static_graph_construction.py

@@ -19,7 +19,7 @@ class StaticGraphConstruction:
 
     # PostEmbedding is costly to instantiate in each StaticGraphConstruction instance.
     post_embedding_builder = PostEmbedding()
-    tag_embedding_model = NextTagEmbeddingTrainer.load_model("../models/tag-emb-1mil.pt", embedding_dim=30, vocab_size=63654, context_length=3)
+    tag_embedding_model = NextTagEmbeddingTrainer.load_model("../models/tag-emb-7_5mil-50d-63653-3.pt", embedding_dim=50, vocab_size=63654, context_length=3)
 
     def __init__(self):
         self._known_tags = {}  # tag_name -> index
@@ -116,7 +116,7 @@ class StaticGraphConstruction:
         if not len(self._known_tags):
             return None
         for tag in self._known_tags:
-            yield StaticGraphConstruction.tag_embedding_model.get_tag_embedding(tag)
+            yield self.tag_embedding_model.get_tag_embedding(tag)
 
     def process_modules(self):
         if not len(self._known_modules):