aurora-model 0.1.0__tar.gz

aurora_model-0.1.0/PKG-INFO

@@ -0,0 +1,29 @@
+Metadata-Version: 2.4
+Name: aurora-model
+Version: 0.1.0
+Summary: Aurora: A multimodal time series prediction model based on Transformers
+Author: Your Name
+Author-email: your_email@example.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Requires-Dist: torch>=2.4.0
+Requires-Dist: torchvision>=0.19.0
+Requires-Dist: transformers>=4.50.0
+Requires-Dist: huggingface_hub>=0.16.0
+Requires-Dist: numpy>=1.21.0
+Requires-Dist: scipy>=1.7.0
+Requires-Dist: einops>=0.8.1
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

aurora_model-0.1.0/aurora/__init__.py

@@ -0,0 +1,91 @@
+# aurora/__init__.py
+import torch
+from typing import Optional
+from huggingface_hub import hf_hub_download
+import os  # ✅ Add os module to check file extension
+
+# Import safetensors for loading .safetensors files
+from safetensors.torch import load_file as safetensors_load_file
+
+# Import your core model and config classes
+from aurora.modeling_aurora import AuroraForPrediction
+from aurora.configuration_aurora import AuroraConfig
+from aurora.utils.path_utils import get_package_file_path
+
+# Hardcode your HF repo info (modify to your actual repo)
+DEFAULT_REPO_ID = "DecisionIntelligence/Aurora"
+# ✅ Update default weights filename to .safetensors (adjust if your file has a different name)
+DEFAULT_WEIGHTS_FILENAME = "model.safetensors"
+
+# Package version
+__version__ = "0.1.0"
+
+
+def load_model(
+        repo_id: str = DEFAULT_REPO_ID,
+        weights_filename: str = DEFAULT_WEIGHTS_FILENAME,
+        cache_dir: Optional[str] = None,
+        force_download: bool = False,
+        device: Optional[str] = None
+) -> AuroraForPrediction:
+    """
+    Load the Aurora model with static files from the package and weights from HF Hub.
+    Supports both .safetensors (recommended) and .bin/.pt/.pth weight formats.
+
+    Args:
+        repo_id: HF Hub repository ID (e.g., "username/repo-name")
+        weights_filename: Name of the weights file in the HF repo (default: model.safetensors)
+        cache_dir: Local directory to cache the downloaded weights (default: HF default cache)
+        force_download: Whether to force re-download the weights even if cached
+        device: Device to load the model onto (default: auto-detect cuda/cpu)
+
+    Returns:
+        Instantiated AuroraForPrediction model with weights loaded
+    """
+    # 1. Load static config file from the installed package
+    config_path = get_package_file_path("config.json")
+    config = AuroraConfig.from_json_file(config_path)
+    print(f"✅ Loaded static config from package: {config_path}")
+
+    # 2. Download only the model weights from HF Hub
+    weights_path = hf_hub_download(
+        repo_id=repo_id,
+        filename=weights_filename,
+        cache_dir=cache_dir,
+        force_download=force_download,
+        resume_download=True
+    )
+    print(f"✅ Downloaded model weights from HF Hub: {weights_path}")
+
+    # 3. Instantiate the model with the static config
+    model = AuroraForPrediction(config)
+
+    # 4. Load weights according to file extension (support .safetensors and .bin/.pt/.pth)
+    file_ext = os.path.splitext(weights_filename)[-1].lower()
+    try:
+        if file_ext == ".safetensors":
+            # ✅ Load .safetensors file (safe, fast, no Pickle risks)
+            weights = safetensors_load_file(weights_path, device="cpu")  # Load to CPU first
+            print("✅ Loaded weights from .safetensors file successfully!")
+        else:
+            # Fallback to torch.load() for traditional PyTorch weight formats
+            weights = torch.load(weights_path, map_location="cpu")
+            print("✅ Loaded weights from traditional PyTorch file successfully!")
+
+        # Load weights into the model
+        model.load_state_dict(weights, strict=False)
+        print("✅ Loaded weights into the model successfully!")
+    except Exception as e:
+        raise RuntimeError(f"Failed to load weights from {weights_path}: {str(e)}") from e
+
+    # 5. Auto-detect and move model to the target device
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    print(f"✅ Model loaded successfully on device: {device}")
+
+    return model
+
+
+# Expose core classes and functions for external use
+__all__ = ["AuroraForPrediction", "AuroraConfig", "load_model"]

aurora_model-0.1.0/aurora/bert_config/config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "BertForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "transformers_version": "4.6.0.dev0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 30522
+}