Upload 4 files

README.md

@@ -4,7 +4,7 @@ emoji: 🫁
 colorFrom: blue
 colorTo: green
 sdk: gradio
-sdk_version: "4.0.0"
+sdk_version: "5.49.1"
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -54,10 +54,11 @@ try:
         ENSEMBLE_CONFIG = json.load(f)
     print("✅ Loaded ensemble_config.json successfully.")
 except FileNotFoundError:
-    print("⚠️ Warning: ensemble_config.json not found. Defaulting to equal weights for ensemble.")
+    print("⚠️ Warning: ensemble_config.json not found. Using default configuration.")
     ENSEMBLE_CONFIG = {"weights": {}}
 except json.JSONDecodeError:
-    print("❌ Error: Could not decode ensemble_config.json. Defaulting to equal weights.")
+    print("❌ Error: Could not decode ensemble_config.json. Using default configuration.")
+    ENSEMBLE_CONFIG = {"weights": {}}
 
 class CBAM(nn.Module):
     """Convolutional Block Attention Module - matches training implementation"""
@@ -146,25 +147,39 @@ def load_ensemble_model(model_repo="calender/Ensemble_C"):
         print("💻 Running locally - checking local files first")
 
     # Map config filenames to local filenames
+    # Local clean files: Model4.pth, Model5.pth, Model6.pth
     config_to_local_map = {
-        "iteration4_best.pth": "Model4.pth",
-        "iteration5_best.pth": "Model5.pth",
-        "iteration6_best.pth": "Model6.pth"
+        "Model4.pth": "Model4.pth",
+        "Model5.pth": "Model5.pth",
+        "Model6.pth": "Model6.pth"
+    }
+
+    # Map config filenames to actual repository filenames
+    # Repository contains clean files: Model4.pth, Model5.pth, Model6.pth
+    # (NOT the checkpoint files Iteration4_BEST.pth which have extra training data)
+    config_to_repo_map = {
+        "Model4.pth": "Model4.pth",
+        "Model5.pth": "Model5.pth",
+        "Model6.pth": "Model6.pth"
     }
 
     # Use filenames from config if available, otherwise default
     config_model_files = list(ENSEMBLE_CONFIG.get("weights", {}).keys())
     if not config_model_files:
         print("⚠️ No model weights in config, using default model files and equal weights.")
-        config_model_files = ["iteration4_best.pth", "iteration5_best.pth", "iteration6_best.pth"]
+        config_model_files = ["Model4.pth", "Model5.pth", "Model6.pth"]
+        print(f"   Expected repository files: {config_model_files}")
+    else:
+        print(f"📋 Config specifies models: {config_model_files}")
+        print(f"🔍 Will download from repository: {[config_to_repo_map.get(f, f) for f in config_model_files]}")
 
-    for hf_filename in config_model_files:
-        weight = ENSEMBLE_CONFIG.get("weights", {}).get(hf_filename, 1.0)
+    for config_filename in config_model_files:
+        weight = ENSEMBLE_CONFIG.get("weights", {}).get(config_filename, 1.0)
 
         if load_locally:
-            local_filename = config_to_local_map.get(hf_filename)
+            local_filename = config_to_local_map.get(config_filename)
             if not local_filename:
-                print(f"❌ Error: No local mapping for '{hf_filename}'. Skipping.")
+                print(f"❌ Error: No local mapping for '{config_filename}'. Skipping.")
                 continue
             model_path = local_filename
             if not os.path.exists(model_path):
@@ -172,14 +187,16 @@ def load_ensemble_model(model_repo="calender/Ensemble_C"):
                 return [], device, []
             print(f"Found local model: {model_path}")
         else:
-            print(f"Downloading {hf_filename} from repo {model_repo}...")
+            # Use the actual filename in the repository
+            repo_filename = config_to_repo_map.get(config_filename, config_filename)
+            print(f"Downloading {repo_filename} from repo {model_repo}...")
             try:
                 # Set HF token for private repo access
                 hf_token = os.environ.get("HF_TOKEN")
                 if hf_token:
                     print(f"✅ Using HF_TOKEN for private repo access")
-                model_path = hf_hub_download(repo_id=model_repo, filename=hf_filename, token=hf_token)
-                print(f"✅ Downloaded {hf_filename} successfully.")
+                model_path = hf_hub_download(repo_id=model_repo, filename=repo_filename, token=hf_token)
+                print(f"✅ Downloaded {repo_filename} successfully.")
 
                 # Basic validation of downloaded file
                 if not model_path.endswith('.pth') and not model_path.endswith('.pt'):
@@ -189,7 +206,8 @@ def load_ensemble_model(model_repo="calender/Ensemble_C"):
                     print(f"❌ Downloaded file not found at {model_path}")
                     continue
             except Exception as e:
-                print(f"❌ Failed to download {hf_filename}: {e}")
+                print(f"❌ Failed to download {repo_filename}: {e}")
+                print(f"   Make sure '{repo_filename}' exists in the '{model_repo}' repository.")
                 continue
 
         # Load state dict with proper error handling
@@ -260,9 +278,9 @@ def load_ensemble_model(model_repo="calender/Ensemble_C"):
 
         try:
             model.load_state_dict(state_dict, strict=False)
-            print(f"✅ Successfully loaded model for '{hf_filename}' with weight {weight:.2f}")
+            print(f"✅ Successfully loaded model for '{config_filename}' with weight {weight:.2f}")
         except Exception as e:
-            print(f"❌ Failed to load state dict for '{hf_filename}': {e}")
+            print(f"❌ Failed to load state dict for '{config_filename}': {e}")
             print("This might indicate architecture mismatch. Consider retraining or using correct model architecture.")
             continue
 
@@ -278,6 +296,9 @@ def load_ensemble_model(model_repo="calender/Ensemble_C"):
     total_weight = sum(model_weights)
     normalized_weights = [w / total_weight for w in model_weights]
     print(f"Ensemble loaded with {len(models)} models. Normalized weights: {[f'{w:.2f}' for w in normalized_weights]}")
+    print(f"📊 Models loaded: {config_model_files}")
+    print(f"🔗 Repository: {model_repo}")
+    print(f"🎯 Ready for Gradio interface!")
 
     return models, device, normalized_weights