modified: modeling.py

modeling.py

@@ -2,7 +2,6 @@ import os
 import torch
 import torch.nn as nn
 from transformers import PreTrainedModel, PretrainedConfig
-from huggingface_hub import snapshot_download
 
 class MVAEConfig(PretrainedConfig):
     model_type = "mvae"
@@ -111,16 +110,7 @@ class MVAE(PreTrainedModel):
             config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
         
         prosoro_type = getattr(config, "prosoro_type", None)
-        print(f"Loading model from {pretrained_model_name_or_path} with prosoro_type={prosoro_type}")
 
-        cached_folder = snapshot_download(pretrained_model_name_or_path)
-        
-        if prosoro_type is not None:
-            sub_path = os.path.join(cached_folder, prosoro_type)
-            print(f"Loading prosoro type from subfolder: {sub_path}")
-            if os.path.isdir(sub_path):
-                pretrained_model_name_or_path = sub_path
-            else:
-                raise ValueError(f"No subfolder found for {prosoro_type}")
+        pretrained_model_name_or_path = pretrained_model_name_or_path + f"/{prosoro_type}"
 
         return super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)