Spaces:
Sleeping
Sleeping
| from huggingface_hub import hf_hub_download | |
| from safetensors.torch import load_model as st_load_model | |
| import os, torch | |
| from agent.sac import Actor, device | |
| # Configuration | |
| REPO_ID = "S-Mithilesh/Energizer-V1" | |
| FILENAME = "model.safetensors" | |
| DATASET_FILENAME = "sample_dataset.pkl" | |
| STATE_DIM = 11 | |
| def load_trained_model(repo_id=REPO_ID, filename=FILENAME, use_local=None): | |
| """ | |
| Loads model weights. | |
| Priority: 1. Forced local (if use_local is set) 2. HF Hub 3. Local safetensors 4. Local .pt | |
| """ | |
| model = Actor(STATE_DIM).to(device) | |
| token = os.environ.get("HF_TOKEN") # Automatic support for private spaces | |
| # 1. Explicit Local Override | |
| if use_local and os.path.exists(use_local): | |
| print(f"Loading forced local weights from {use_local}...") | |
| try: | |
| st_load_model(model, use_local) | |
| print("[SUCCESS] Successfully loaded forced local weights.") | |
| return model.eval() | |
| except Exception as e: | |
| print(f"[ERROR] Failed to load forced local weights: {e}") | |
| # 2. Try Hugging Face Hub (Preferred) | |
| try: | |
| print(f"Fetching weights from HF Hub: {repo_id}/{filename}...") | |
| # Download file from HF Hub (cached locally) | |
| path = hf_hub_download(repo_id=repo_id, filename=filename, token=token) | |
| st_load_model(model, path) | |
| print("[SUCCESS] Model weights loaded successfully from HF Hub.") | |
| return model.eval() | |
| except Exception as e: | |
| print(f"[WARNING] Hub fetch failed: {e}") | |
| # 3. Fallback to local safetensors (if exists in repo) | |
| local_path = "models/model.safetensors" | |
| if os.path.exists(local_path): | |
| print(f"[FALLBACK] Loading local weights from {local_path}...") | |
| try: | |
| st_load_model(model, local_path) | |
| print("[SUCCESS] Loaded local safetensors fallback.") | |
| return model.eval() | |
| except Exception as e: | |
| print(f"[WARNING] Local safetensors load failed: {e}") | |
| # 4. Final fallback for original .pt file | |
| pt_path = "models/policy_sac.pt" | |
| if os.path.exists(pt_path): | |
| print(f"[FALLBACK] Loading local .pt checkpoint...") | |
| try: | |
| ckpt = torch.load(pt_path, map_location=device, weights_only=False) | |
| sd = ckpt['actor_state'] if isinstance(ckpt, dict) and 'actor_state' in ckpt else ckpt | |
| model.load_state_dict(sd) | |
| print("[SUCCESS] Loaded local .pt fallback.") | |
| except Exception as e: | |
| print(f"[ERROR] Local .pt load failed: {e}") | |
| else: | |
| print("[ERROR] No weights found. Model will have random performance.") | |
| model.eval() | |
| return model | |
| def ensure_dataset_exists(repo_id=REPO_ID, filename=DATASET_FILENAME, target_dir="data"): | |
| """ | |
| Ensures the sample dataset is available locally. | |
| Priority: 1. Local data/ folder 2. HF Hub download | |
| """ | |
| target_path = os.path.join(target_dir, filename) | |
| if os.path.exists(target_path): | |
| return target_path | |
| token = os.environ.get("HF_TOKEN") | |
| try: | |
| print(f"Fetching sample dataset from HF Hub: {repo_id}/{filename}") | |
| os.makedirs(target_dir, exist_ok=True) | |
| # Download and cache | |
| path = hf_hub_download(repo_id=repo_id, filename=filename, token=token) | |
| # Move/Link to the expected data directory for environment consistency | |
| import shutil | |
| shutil.copy(path, target_path) | |
| print(f"[SUCCESS] Sample dataset localized to {target_path}") | |
| return target_path | |
| except Exception as e: | |
| print(f"[WARNING] Could not fetch sample dataset from Hub: {e}") | |
| return None | |
| if __name__ == "__main__": | |
| # Test loading | |
| m = load_trained_model() | |
| d = ensure_dataset_exists() | |
| print("Test complete.") | |