Fix model to load weights

fromage/models.py

@@ -634,21 +634,21 @@ def load_fromage(embeddings_dir: str, model_args_path: str, model_ckpt_path: str
   ret_token_idx = tokenizer('[RET]', add_special_tokens=False).input_ids
   assert len(ret_token_idx) == 1, ret_token_idx
   model_kwargs['retrieval_token_idx'] = ret_token_idx[0]
-  model_kwargs['opt_version'] = 'facebook/opt-125m'
-  model_kwargs['visual_encoder'] = 'openai/clip-vit-base-patch32'
+  # model_kwargs['opt_version'] = 'facebook/opt-125m'
+  # model_kwargs['visual_encoder'] = 'openai/clip-vit-base-patch32'
   args = namedtuple('args', model_kwargs)(**model_kwargs)
 
   # Initialize model for inference.
   model = Fromage(tokenizer, args, path_array=path_array, emb_matrix=emb_matrix)
   model = model.eval()
-  # model = model.bfloat16()
-  # model = model.cuda()
+  model = model.bfloat16()
+  model = model.cuda()
 
   # Load pretrained linear mappings and [RET] embeddings.
-  # checkpoint = torch.load(model_ckpt_path)
-  # model.load_state_dict(checkpoint['state_dict'], strict=False)
-  # with torch.no_grad():
-  #     model.model.input_embeddings.weight[model.model.retrieval_token_idx, :].copy_(checkpoint['state_dict']['ret_input_embeddings.weight'].cpu().detach())
+  checkpoint = torch.load(model_ckpt_path)
+  model.load_state_dict(checkpoint['state_dict'], strict=False)
+  with torch.no_grad():
+      model.model.input_embeddings.weight[model.model.retrieval_token_idx, :].copy_(checkpoint['state_dict']['ret_input_embeddings.weight'].cpu().detach())
 
   logit_scale = model.model.logit_scale.exp()
   emb_matrix = torch.tensor(emb_matrix, dtype=logit_scale.dtype).to(logit_scale.device)