initial commit

app.py

@@ -70,26 +70,29 @@ class MultimodalRiskBehaviorModel(nn.Module):
     @classmethod
     def from_pretrained(cls, load_directory, map_location=None):
         if os.path.exists(load_directory):
-            # Local directory: load configuration from file
-            with open(os.path.join(load_directory, 'config.json'), 'r') as f:
-                config_dict = json.load(f)
-        else:
-            # Hugging Face Hub: load configuration using AutoConfig
-            config = AutoConfig.from_pretrained(load_directory)
-            config_dict = {
-                "text_model_name": config.name_or_path,
-                "hidden_dim": config.hidden_size if hasattr(config, 'hidden_size') else 512  # Default to 512 if not available
-            }
-        
-        model = cls(text_model_name=config_dict["text_model_name"], hidden_dim=config_dict["hidden_dim"])
+            # Handling local paths
+            config_path = os.path.join(load_directory, 'config.json')
+            state_dict_path = os.path.join(load_directory, 'pytorch_model.bin')
 
-        state_dict_path = os.path.join(load_directory, 'pytorch_model.bin')
-        if map_location:
+            if not os.path.exists(config_path):
+                raise FileNotFoundError(f"{config_path} not found.")
+            if not os.path.exists(state_dict_path):
+                raise FileNotFoundError(f"{state_dict_path} not found.")
+
+            with open(config_path, 'r') as f:
+                config_dict = json.load(f)
+            model = cls(text_model_name=config_dict["text_model_name"], hidden_dim=config_dict["hidden_dim"])
             state_dict = torch.load(state_dict_path, map_location=map_location)
+            model.load_state_dict(state_dict)
+        
         else:
-            state_dict = torch.load(state_dict_path)
+            # Handling Hugging Face Hub paths
+            config = AutoConfig.from_pretrained(load_directory)
 
-        model.load_state_dict(state_dict)
+            # Use AutoModel to load a pre-trained model with weights
+            hf_model = AutoModel.from_pretrained(load_directory, config=config)
+            model = cls(text_model_name=config.name_or_path, hidden_dim=hf_model.config.hidden_size)
+            model.text_model = hf_model
 
         return model