Fix loading

fromage/models.py

@@ -594,17 +594,15 @@ class Fromage(nn.Module):
     return return_outputs
 
 
-def load_fromage(embeddings_dir: str, args_path: str, ckpt_path: str) -> Fromage:
-  model_args_path = os.path.join(model_dir, 'model_args.json')
-  model_ckpt_path = os.path.join(model_dir, 'pretrained_ckpt.pth.tar')
-  embs_paths = [s for s in glob.glob(os.path.join(model_dir, 'cc3m_embeddings*.pkl'))]
+def load_fromage(embeddings_dir: str, model_args_path: str, model_ckpt_path: str) -> Fromage:
+  embs_paths = [s for s in glob.glob(os.path.join(embeddings_dir, 'cc3m_embeddings*.pkl'))]
 
   if not os.path.exists(model_args_path):
-    raise ValueError(f'model_args.json does not exist in {model_dir}.')
+    raise ValueError(f'model_args.json does not exist at {model_args_path}.')
   if not os.path.exists(model_ckpt_path):
-    raise ValueError(f'pretrained_ckpt.pth.tar does not exist in {model_dir}.')
+    raise ValueError(f'pretrained_ckpt.pth.tar does not exist at {model_ckpt_path}.')
   if len(embs_paths) == 0:
-    raise ValueError(f'cc3m_embeddings_*.pkl files do not exist in {model_dir}.')
+    raise ValueError(f'cc3m_embeddings_*.pkl files do not exist in {embeddings_dir}.')
 
   # Load embeddings.
   # Construct embedding matrix for nearest neighbor lookup.