Add huggingface support.

Author: zmgong

bioscanclip/config/global_config.yaml

@@ -62,4 +62,6 @@ general_fine_tune_setting:
   epoch: 15
   batch_size: 200
 
+hf_repo_id: bioscan-ml/clibd
+
 default_seed: 42

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_seed_42.yaml

@@ -15,6 +15,7 @@ dna:
 model_output_name: image_dna_4gpu
 evaluation_period: 1
 ckpt_path: ${project_root_path}/ckpt/bioscan_clip/final_experiments/image_dna_4gpu_50epoch/best.pth
+hf_model_name: ckpt/bioscan_clip/final_experiments/image_dna_4gpu_50epoch/best.pth
 output_dim: 768
 port: 29532
 

bioscanclip/config/model_config/for_bioscan_1m/final_experiments/image_dna_text_seed_42.yaml

@@ -17,6 +17,7 @@ language:
 model_output_name: image_dna_text_4gpu
 evaluation_period: 1
 ckpt_path: ${project_root_path}/ckpt/bioscan_clip/final_experiments/image_dna_text_4gpu_50epoch/best.pth
+hf_model_name: ckpt/bioscan_clip/final_experiments/image_dna_text_4gpu_50epoch/best.pth
 output_dim: 768
 port: 29531
 

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_seed_42.yaml

@@ -15,6 +15,7 @@ dna:
 model_output_name: image_dna_4gpu
 evaluation_period: 1
 ckpt_path: ${project_root_path}/ckpt/bioscan_clip/new_5M_training/trained_with_5M_image_dna/best.pth
+hf_model_name: ckpt/bioscan_clip/new_5M_training/trained_with_5M_image_dna/best.pth
 output_dim: 768
 port: 29531
 

bioscanclip/config/model_config/for_bioscan_5m/final_experiments/image_dna_text_seed_42.yaml

@@ -18,6 +18,7 @@ language:
 model_output_name: image_dna_text_4gpu
 evaluation_period: 1
 ckpt_path: ${project_root_path}/ckpt/bioscan_clip/new_5M_training/trained_with_5M_image_dna_text/best.pth
+hf_model_name: ckpt/bioscan_clip/new_5M_training/trained_with_5M_image_dna_text/best.pth
 output_dim: 768
 port: 29531
 

scripts/inference_and_eval.py

@@ -26,6 +26,7 @@
     make_prediction,
     All_TYPE_OF_FEATURES_OF_KEY,
 )
+from huggingface_hub import hf_hub_download
 
 PLOT_FOLDER = "html_plots"
 RETRIEVAL_FOLDER = "image_retrieval"
@@ -595,12 +596,23 @@ def main(args: DictConfig) -> None:
 
         if hasattr(args.model_config, "load_ckpt") and args.model_config.load_ckpt is False:
             pass
+        # elif os.path.exists(args.model_config.ckpt_path):
+        #     checkpoint = torch.load(args.model_config.ckpt_path, map_location="cuda:0")
+        #     print(f"Loading model from {args.model_config.ckpt_path}")
+        #     print()
+        #     model.load_state_dict(checkpoint)
+        elif hasattr(args.model_config, "hf_model_name"):
+            checkpoint_path = hf_hub_download(
+                repo_id=args.hf_repo_id,
+                filename=args.model_config.hf_model_name,
+            )
+            checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
+            model.load_state_dict(checkpoint)
+            print(f"Loading model from {args.hf_repo_id}/{args.model_config.hf_model_name}")
         else:
-            checkpoint = torch.load(args.model_config.ckpt_path, map_location="cuda:0")
-            print(f"Loading model from {args.model_config.ckpt_path}")
-            print()
+            raise ValueError("No checkpoint found. Please specify a valid checkpoint path.")
+
 
-            model.load_state_dict(checkpoint)
 
         # Load data
         # args.model_config.batch_size = 24