Upload 5 files

__init__.py

@@ -0,0 +1,4 @@
+from .configuration import WeatherSequenceConfig
+from .modeling import WeatherSequenceModel, WeatherSequenceOutput
+
+__all__ = ["WeatherModelConfig", "WeatherForcastModel", "WeatherModelOutput"]

config.json

@@ -0,0 +1,82 @@
+{
+  "architectures": [
+    "WeatherSequenceModel"
+  ],
+  "distill_teacher_head_dim": 416,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "encoder_type": "lstm",
+  "hidden_dim": 128,
+  "input_dim": 22,
+  "location_emb_dim": 32,
+  "model_type": "weather_sequence",
+  "num_layers": 3,
+  "num_locations": 82,
+  "num_predict": 12,
+  "num_weather_classes": 7,
+  "rain_pos_weight": 6.547722074664306,
+  "seq_len": 72,
+  "target_norms": {
+    "apparent": {
+      "mean": 16.420160986060196,
+      "std": 12.332221212388726,
+      "transform": "raw"
+    },
+    "cloud_cover": {
+      "mean": 51.71581237675868,
+      "std": 42.028595137718646,
+      "transform": "raw"
+    },
+    "humidity": {
+      "mean": 69.18635597409919,
+      "std": 21.524024234467674,
+      "transform": "raw"
+    },
+    "precip": {
+      "mean": 0.05684371705656333,
+      "std": 0.22550783339649325,
+      "transform": "log1p"
+    },
+    "sea_level_pressure": {
+      "mean": 1014.5679196119568,
+      "std": 7.459071118489876,
+      "transform": "raw"
+    },
+    "surface_pressure": {
+      "mean": 963.4371186423618,
+      "std": 86.39541603431283,
+      "transform": "raw"
+    },
+    "temp": {
+      "mean": 16.863813962719767,
+      "std": 9.92933797011761,
+      "transform": "raw"
+    },
+    "wind": {
+      "mean": 9.952284635073887,
+      "std": 6.676419945197847,
+      "transform": "raw"
+    },
+    "wind_dir_cos": {
+      "mean": -0.03349104536155545,
+      "std": 0.70513783656826,
+      "transform": "raw"
+    },
+    "wind_dir_sin": {
+      "mean": -0.002100160488024665,
+      "std": 0.7082757736110976,
+      "transform": "raw"
+    }
+  },
+  "transformers_version": "5.5.0",
+  "use_cache": false,
+  "weather_class_weights": [
+    0.23165243864059448,
+    0.1921183317899704,
+    1.7011765241622925,
+    0.4377932548522949,
+    0.8518651723861694,
+    1.6312177181243896,
+    1.9541765451431274
+  ]
+}

configuration.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+try:
+    from transformers import PretrainedConfig
+except Exception:  # pragma: no cover - lets the file import in minimal environments
+    class PretrainedConfig:  # type: ignore
+        model_type = "custom"
+
+        def __init__(self, **kwargs):
+            for k, v in kwargs.items():
+                setattr(self, k, v)
+
+
+class WeatherModelConfig(PretrainedConfig):
+
+    model_type = "mwm"
+
+    def __init__(
+        self,
+        input_dim: Optional[int] = None,
+        seq_len: int = 72,
+        num_predict: int = 12,
+        hidden_dim: int = 128,
+        num_layers: int = 3,
+        dropout: float = 0.1,
+        encoder_type: str = "lstm",
+        num_locations: int = 82,
+        location_emb_dim: int = 32,
+        num_weather_classes: int = 7,
+        rain_pos_weight: float = 1.0,
+        weather_class_weights: Optional[list[float]] = None,
+        target_norms: Optional[Dict[str, Dict[str, float]]] = None,
+        distill_teacher_head_dim: int = 416,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.input_dim = input_dim
+        self.seq_len = seq_len
+        self.num_predict = num_predict
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.encoder_type = encoder_type
+        self.num_locations = num_locations
+        self.location_emb_dim = location_emb_dim
+        self.num_weather_classes = num_weather_classes
+        self.rain_pos_weight = rain_pos_weight
+        self.weather_class_weights = weather_class_weights
+        self.target_norms = target_norms or {}
+        self.distill_teacher_head_dim = int(distill_teacher_head_dim)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca32ca6acc93e3d503343366a6cc277b723d82ad1cce2815dc7fe761772a0748
+size 1788776

modeling.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from transformers import PreTrainedModel
+    from transformers.modeling_outputs import ModelOutput
+except Exception: 
+    class PreTrainedModel(nn.Module): 
+        config_class = None
+        base_model_prefix = ""
+        main_input_name = "input_ids"
+
+        def __init__(self, config):
+            super().__init__()
+            self.config = config
+
+    class ModelOutput(dict):  # type: ignore
+        pass
+
+from .configuration import WeatherModelConfig
+
+CONTINUOUS_TARGET_ORDER = [
+    "temp",
+    "humidity",
+    "apparent",
+    "precip",
+    "sea_level_pressure",
+    "surface_pressure",
+    "cloud_cover",
+    "wind",
+    "wind_dir_sin",
+    "wind_dir_cos",
+]
+
+CONTINUOUS_TARGET_SPECS = {
+    "temp": {"loss_weight": 1.0, "transform": "raw"},
+    "humidity": {"loss_weight": 1.0, "transform": "raw"},
+    "apparent": {"loss_weight": 0.8, "transform": "raw"},
+    "precip": {"loss_weight": 0.9, "transform": "log1p"},
+    "sea_level_pressure": {"loss_weight": 0.6, "transform": "raw"},
+    "surface_pressure": {"loss_weight": 0.4, "transform": "raw"},
+    "cloud_cover": {"loss_weight": 0.4, "transform": "raw"},
+    "wind": {"loss_weight": 0.6, "transform": "raw"},
+    "wind_dir_sin": {"loss_weight": 0.55, "transform": "raw"},
+    "wind_dir_cos": {"loss_weight": 0.55, "transform": "raw"},
+}
+
+
+@dataclass
+class WeatherModelOutput(ModelOutput):
+    loss: Optional[torch.Tensor] = None
+    logits: Optional[Tuple[torch.Tensor, ...]] = None
+    head_repr: Optional[torch.Tensor] = None
+    norm_preds: Optional[Dict[str, torch.Tensor]] = None
+    raw_preds: Optional[Dict[str, torch.Tensor]] = None
+    distill_head_repr: Optional[torch.Tensor] = None
+
+
+class WeatherForcastModel(PreTrainedModel):
+
+    config_class = WeatherModelConfig
+    base_model_prefix = "weather_sequence"
+    main_input_name = "X"
+
+    # Newer Transformers versions may create auto_map entries from these registrations.
+    _tied_weights_keys: list[str] = []
+
+    def __init__(self, config: WeatherModelConfig):
+        super().__init__(config)
+
+        self.encoder_type = str(getattr(config, "encoder_type", "lstm")).lower()
+        self.hidden_dim = int(config.hidden_dim)
+        self.seq_len = int(config.seq_len)
+        self.num_predict = int(config.num_predict)
+        self.num_weather_classes = int(config.num_weather_classes)
+
+        if config.input_dim is None:
+            raise ValueError("WeatherModelConfig.input_dim must be set")
+
+        self.location_embedding = nn.Embedding(max(1, int(config.num_locations)), int(config.location_emb_dim))
+
+        if config.weather_class_weights is not None:
+            self.register_buffer(
+                "weather_class_weights",
+                torch.tensor(config.weather_class_weights, dtype=torch.float32),
+                persistent=False,
+            )
+        else:
+            self.weather_class_weights = None
+
+        self.register_buffer(
+            "rain_pos_weight",
+            torch.tensor(float(config.rain_pos_weight), dtype=torch.float32),
+            persistent=False,
+        )
+
+        self.target_norm_meta: Dict[str, Dict[str, Any]] = {}
+        for name in CONTINUOUS_TARGET_ORDER:
+            spec = dict(config.target_norms.get(name, {}))
+            mean = float(spec.get("mean", 0.0))
+            std = max(float(spec.get("std", 1.0)), 1e-6)
+            transform = str(spec.get("transform", CONTINUOUS_TARGET_SPECS[name]["transform"]))
+            self.register_buffer(f"{name}_mean", torch.tensor(mean, dtype=torch.float32), persistent=False)
+            self.register_buffer(f"{name}_std", torch.tensor(std, dtype=torch.float32), persistent=False)
+            self.target_norm_meta[name] = {"transform": transform}
+
+        if self.encoder_type == "lstm":
+            self.encoder = nn.LSTM(
+                input_size=int(config.input_dim),
+                hidden_size=self.hidden_dim,
+                num_layers=int(config.num_layers),
+                batch_first=True,
+                dropout=float(config.dropout) if int(config.num_layers) > 1 else 0.0,
+                bidirectional=False,
+            )
+        elif self.encoder_type == "transformer":
+            self.input_proj = nn.Linear(int(config.input_dim), self.hidden_dim)
+            self.pos_encoding = nn.Parameter(torch.randn(1, int(config.seq_len), self.hidden_dim) * 0.1)
+            encoder_layer = nn.TransformerEncoderLayer(
+                d_model=self.hidden_dim,
+                nhead=4,
+                dropout=float(config.dropout),
+                batch_first=True,
+            )
+            self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=int(config.num_layers))
+        else:
+            raise ValueError(f"Unknown encoder_type: {self.encoder_type}")
+
+        self.head_dim = self.hidden_dim + int(config.location_emb_dim)
+        self.head_norm = nn.LayerNorm(self.head_dim)
+        self.head_dropout = nn.Dropout(float(config.dropout))
+
+        self.reg_heads = nn.ModuleDict({name: nn.Linear(self.head_dim, self.num_predict) for name in CONTINUOUS_TARGET_ORDER})
+        self.fc_rain = nn.Linear(self.head_dim, self.num_predict)
+        self.fc_weather = nn.Linear(self.head_dim, self.num_predict * self.num_weather_classes)
+
+        teacher_head_dim = int(getattr(config, "distill_teacher_head_dim", 0))
+        if teacher_head_dim > 0 and teacher_head_dim != self.head_dim:
+            self.distill_proj = nn.Linear(self.head_dim, teacher_head_dim, bias=False)
+        else:
+            self.distill_proj = None
+
+        self.post_init()
+
+    @staticmethod
+    def _masked_mean(x: torch.Tensor) -> torch.Tensor:
+        mask = (x.abs().sum(dim=-1) > 0).float().unsqueeze(-1)
+        summed = (x * mask).sum(dim=1)
+        denom = mask.sum(dim=1).clamp(min=1.0)
+        return summed / denom
+
+    def _target_mean_std(self, name: str) -> Tuple[torch.Tensor, torch.Tensor]:
+        return getattr(self, f"{name}_mean"), getattr(self, f"{name}_std")
+
+    def _encode_target(self, name: str, target: torch.Tensor) -> torch.Tensor:
+        transform = self.target_norm_meta[name]["transform"]
+        target = target.to(dtype=torch.float32)
+        if transform == "log1p":
+            target = torch.log1p(torch.clamp(target, min=0.0))
+        mean, std = self._target_mean_std(name)
+        return (target - mean.to(target.device)) / std.to(target.device)
+
+    def _decode_prediction(self, name: str, pred_norm: torch.Tensor) -> torch.Tensor:
+        transform = self.target_norm_meta[name]["transform"]
+        mean, std = self._target_mean_std(name)
+        raw = pred_norm * std.to(pred_norm.device) + mean.to(pred_norm.device)
+        if transform == "log1p":
+            raw = torch.expm1(raw).clamp(min=0.0)
+        return raw
+
+    def forward(
+        self,
+        X: torch.Tensor,
+        location_id: Optional[torch.Tensor] = None,
+        temp_target: Optional[torch.Tensor] = None,
+        humidity_target: Optional[torch.Tensor] = None,
+        apparent_target: Optional[torch.Tensor] = None,
+        precip_target: Optional[torch.Tensor] = None,
+        sea_level_pressure_target: Optional[torch.Tensor] = None,
+        surface_pressure_target: Optional[torch.Tensor] = None,
+        cloud_cover_target: Optional[torch.Tensor] = None,
+        wind_target: Optional[torch.Tensor] = None,
+        wind_dir_sin_target: Optional[torch.Tensor] = None,
+        wind_dir_cos_target: Optional[torch.Tensor] = None,
+        rain_target: Optional[torch.Tensor] = None,
+        weather_target: Optional[torch.Tensor] = None,
+        return_repr: bool = False,
+        **kwargs: Any,
+    ) -> WeatherModelOutput:
+        if location_id is None:
+            location_id = torch.zeros(X.size(0), dtype=torch.long, device=X.device)
+
+        if self.encoder_type == "lstm":
+            _, (h, _) = self.encoder(X)
+            seq_repr = h[-1]
+        else:
+            z = self.input_proj(X) + self.pos_encoding[:, : X.size(1), :]
+            out = self.encoder(z)
+            seq_repr = self._masked_mean(out)
+
+        loc_emb = self.location_embedding(location_id)
+        head_repr = self.head_norm(torch.cat([seq_repr, loc_emb], dim=1))
+        h = self.head_dropout(head_repr)
+
+        raw_preds: Dict[str, torch.Tensor] = {}
+        norm_preds: Dict[str, torch.Tensor] = {}
+        for name in CONTINUOUS_TARGET_ORDER:
+            norm_pred = self.reg_heads[name](h)
+            norm_preds[name] = norm_pred
+            raw_preds[name] = self._decode_prediction(name, norm_pred)
+
+        rain_logit = self.fc_rain(h)
+        weather_logits = self.fc_weather(h).view(-1, self.num_predict, self.num_weather_classes)
+
+        loss = None
+        if temp_target is not None:
+            targets = {
+                "temp": temp_target,
+                "humidity": humidity_target,
+                "apparent": apparent_target,
+                "precip": precip_target,
+                "sea_level_pressure": sea_level_pressure_target,
+                "surface_pressure": surface_pressure_target,
+                "cloud_cover": cloud_cover_target,
+                "wind": wind_target,
+                "wind_dir_sin": wind_dir_sin_target,
+                "wind_dir_cos": wind_dir_cos_target,
+            }
+
+            loss_terms = []
+            for name, target in targets.items():
+                if target is None:
+                    continue
+                target_norm = self._encode_target(name, target.to(h.device))
+                pred_norm = norm_preds[name].to(target_norm.dtype)
+                loss_terms.append(
+                    F.smooth_l1_loss(pred_norm, target_norm) * float(CONTINUOUS_TARGET_SPECS[name]["loss_weight"])
+                )
+
+            if rain_target is not None:
+                rain_target = rain_target.to(rain_logit.dtype)
+                rain_loss = F.binary_cross_entropy_with_logits(
+                    rain_logit,
+                    rain_target,
+                    pos_weight=self.rain_pos_weight.to(rain_logit.device),
+                )
+                loss_terms.append(0.7 * rain_loss)
+
+            if weather_target is not None:
+                weather_loss = F.cross_entropy(
+                    weather_logits.reshape(-1, self.num_weather_classes),
+                    weather_target.long().reshape(-1),
+                    weight=self.weather_class_weights,
+                    label_smoothing=0.0,
+                )
+                loss_terms.append(0.9 * weather_loss)
+
+            loss = sum(loss_terms) if loss_terms else None
+
+        logits = (
+            raw_preds["temp"],
+            raw_preds["humidity"],
+            raw_preds["apparent"],
+            raw_preds["precip"],
+            raw_preds["sea_level_pressure"],
+            raw_preds["surface_pressure"],
+            raw_preds["cloud_cover"],
+            raw_preds["wind"],
+            raw_preds["wind_dir_sin"],
+            raw_preds["wind_dir_cos"],
+            rain_logit,
+            weather_logits,
+        )
+
+        output = WeatherModelOutput(
+            loss=loss,
+            logits=logits,
+            head_repr=head_repr if return_repr else None,
+            norm_preds=norm_preds if return_repr else None,
+            raw_preds=raw_preds if return_repr else None,
+            distill_head_repr=(self.distill_proj(head_repr) if self.distill_proj is not None else head_repr) if return_repr else None,
+        )
+        return output
+
+
+# Make the repo usable with AutoConfig/AutoModel when loaded from the Hub.
+try:  # pragma: no cover
+    WeatherModelConfig.register_for_auto_class()
+except Exception:
+    pass
+
+try:  # pragma: no cover
+    WeatherForcastModel.register_for_auto_class("AutoModel")
+except Exception:
+    pass