clean up annoy warning in predict.py

raubitsj · raubitsj · commit 386a402e3fc0 · 2019-09-24T23:40:56.000-07:00
diff --git a/src/predict.py b/src/predict.py
@@ -111,7 +111,7 @@ def query_model(query, model, indices, language, topk=100):
         indexes = [{'code_tokens': d['function_tokens'], 'language': d['language']} for d in tqdm(definitions)]
         code_representations = model.get_code_representations(indexes)
 
-        indices = AnnoyIndex(code_representations[0].shape[0])
+        indices = AnnoyIndex(code_representations[0].shape[0], 'angular')
         for index, vector in tqdm(enumerate(code_representations)):
             if vector is not None:
                 indices.add_item(index, vector)
