index edge2vec converts names for agatha training

JSybrandt · May 22, 2020 · c1c7b1f · c1c7b1f
1 parent 82b8588
commit c1c7b1f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/agatha/ml/hypothesis_predictor/hypothesis_predictor.py b/agatha/ml/hypothesis_predictor/hypothesis_predictor.py
@@ -457,6 +457,11 @@ def add_argparse_args(parser:ArgumentParser)->ArgumentParser:
     parser.add_argument("--warmup-steps", type=int)
     parser.add_argument("--weight-decay", type=float)
     parser.add_argument("--disable-cache", action="store_true")
+    parser.add_argument(
+        "--simple",
+        help="If set, ignore graph and use a simpler model architecture.",
+        action="store_true"
+    )
     return parser
 
   def init_ddp_connection(

diff --git a/tools/py_scripts/edge2vec_conversion/index_edge2vec_embeddings.py b/tools/py_scripts/edge2vec_conversion/index_edge2vec_embeddings.py
@@ -22,6 +22,7 @@
 import pickle
 from typing import Iterable, List, Tuple
 from tqdm import tqdm
+from agatha.ml.hypothesis_predictor.predicate_util import clean_coded_term
 
 
 def iterate_vectors(vector_text_path:Path)->Iterable[Tuple[str, List[float]]]:
@@ -61,7 +62,7 @@ def main(
   idx2node = {i: n for n, i in node2idx.items()}
 
   node2vec = {
-      idx2node[idx]: vec
+      clean_coded_term(idx2node[idx]): vec
       for idx, vec
       in tqdm(iterate_vectors(input_vector_text_path))
   }