import json import numpy as np import torch import gradio as gr from transformers import AutoConfig, PreTrainedModel, PretrainedConfig from safetensors.torch import load_file from huggingface_hub import hf_hub_download from typing import Optional, Dict, List import math import torch.nn as nn MODEL_ID = "SupraLabs/SupraWeather1.5-Small" # ── Label / feature maps ───────────────────────────────────────────────────── PHENOMENA = [ 'clear','partly_cloudy','cloudy','overcast','mist','fog','dense_fog', 'light_rain','rain','heavy_rain','torrential_rain','thunderstorm', 'severe_thunderstorm','snow','heavy_snow','blizzard','freezing_rain', 'ice_storm','soft_hail','sleet','cold_front','heat_wave','cold_wave', 'windstorm','dust_storm', ] NUM_CLASSES = len(PHENOMENA) LABEL2ID = {p: i for i, p in enumerate(PHENOMENA)} ID2LABEL = {i: p for i, p in enumerate(PHENOMENA)} CONTINUOUS_FEATURES = [ 'temperature','humidity','pressure','pressure_trend','wind_speed', 'altitude','cloud_cover','visibility','solar_radiation','dew_point', 'heat_index','wind_chill','storm_index','rain_potential','instability_index', ] CAT_CARDINALITIES = [8, 12, 6, 6, 5] WIND_DIR_MAP = {'north':0,'northeast':1,'east':2,'southeast':3,'south':4,'southwest':5,'west':6,'northwest':7} AIR_MASS_MAP = {'polar':0,'arctic':1,'continental':2,'maritime':3,'tropical':4,'equatorial':5} CLIMATE_ZONE_MAP = {'polar':0,'subarctic':1,'temperate':2,'mediterranean':3,'subtropical':4,'tropical':5} TERRAIN_MAP = {'flat':0,'coastal':1,'valley':2,'hills':3,'mountain':4} # ── Model architecture (must match training) ────────────────────────────────── class SupraWeatherConfig(PretrainedConfig): model_type = 'supra_weather_ft' def __init__(self, num_continuous=15, cat_cardinalities=None, d_token=256, n_blocks=4, n_heads=8, ffn_factor=1.333, attn_dropout=0.15, ffn_dropout=0.10, residual_dropout=0.0, num_labels=25, means=None, stds=None, continuous_features=None, categorical_features=None, label2id=None, id2label=None, **kwargs): super().__init__(**kwargs) self.num_continuous = num_continuous self.cat_cardinalities = cat_cardinalities or CAT_CARDINALITIES self.d_token = d_token self.n_blocks = n_blocks self.n_heads = n_heads self.ffn_factor = ffn_factor self.attn_dropout = attn_dropout self.ffn_dropout = ffn_dropout self.residual_dropout = residual_dropout self.num_labels = num_labels self.means = means or {} self.stds = stds or {} self.continuous_features = continuous_features or CONTINUOUS_FEATURES self.categorical_features = categorical_features or [] self.label2id = label2id or LABEL2ID self.id2label = {int(k): v for k, v in id2label.items()} if id2label else ID2LABEL class NumericalTokenizer(nn.Module): def __init__(self, n, d): super().__init__() self.W = nn.Parameter(torch.empty(n, d)) self.b = nn.Parameter(torch.zeros(n, d)) nn.init.kaiming_uniform_(self.W, a=math.sqrt(5)) def forward(self, x): return x.unsqueeze(-1) * self.W + self.b class CategoricalTokenizer(nn.Module): def __init__(self, cardinalities, d): super().__init__() self.embs = nn.ModuleList([nn.Embedding(c + 1, d) for c in cardinalities]) def forward(self, cats): return torch.stack([e(c) for e, c in zip(self.embs, cats)], dim=1) class FTBlock(nn.Module): def __init__(self, cfg): super().__init__() D = cfg.d_token D_ffn = max(int(D * cfg.ffn_factor), 4) self.ln1 = nn.LayerNorm(D) self.attn = nn.MultiheadAttention(D, cfg.n_heads, dropout=cfg.attn_dropout, batch_first=True) self.ln2 = nn.LayerNorm(D) self.ffn = nn.Sequential(nn.Linear(D, D_ffn), nn.GELU(), nn.Dropout(cfg.ffn_dropout), nn.Linear(D_ffn, D)) self.drop = nn.Dropout(cfg.residual_dropout) def forward(self, x): h, _ = self.attn(*([self.ln1(x)]*3)) x = x + self.drop(h) x = x + self.drop(self.ffn(self.ln2(x))) return x class SupraWeatherModel(PreTrainedModel): config_class = SupraWeatherConfig @classmethod def _can_set_experts_implementation(cls): return False def __init__(self, cfg): super().__init__(cfg) D = cfg.d_token self.num_tok = NumericalTokenizer(cfg.num_continuous, D) self.cat_tok = CategoricalTokenizer(cfg.cat_cardinalities, D) self.cls = nn.Parameter(torch.zeros(1, 1, D)) self.blocks = nn.ModuleList([FTBlock(cfg) for _ in range(cfg.n_blocks)]) self.ln_out = nn.LayerNorm(D) self.head = nn.Sequential(nn.Linear(D, D), nn.GELU(), nn.Dropout(cfg.ffn_dropout), nn.Linear(D, cfg.num_labels)) self.post_init() def forward(self, continuous, wind_direction, month, air_mass, climate_zone, terrain_type, labels=None, **kwargs): B = continuous.size(0) x = torch.cat([self.cls.expand(B,-1,-1), self.num_tok(continuous), self.cat_tok([wind_direction, month, air_mass, climate_zone, terrain_type])], dim=1) for blk in self.blocks: x = blk(x) logits = self.head(self.ln_out(x[:, 0])) loss = nn.CrossEntropyLoss()(logits, labels) if labels is not None else None return {'loss': loss, 'logits': logits} if loss is not None else {'logits': logits} # ── Load model once ─────────────────────────────────────────────────────────── AutoConfig.register('supra_weather_ft', SupraWeatherConfig) cfg = SupraWeatherConfig.from_pretrained(MODEL_ID) model = SupraWeatherModel(cfg) sd = load_file(hf_hub_download(MODEL_ID, 'model.safetensors')) model.load_state_dict(sd) model.eval() MEANS = cfg.means STDS = cfg.stds # ── Inference ───────────────────────────────────────────────────────────────── def derive_and_predict(temperature, humidity, pressure, pressure_trend, wind_speed, altitude, cloud_cover, visibility, solar_radiation, wind_direction, month, air_mass, climate_zone, terrain_type): dew_point = temperature - (100 - humidity) / 5.0 heat_index = temperature if temperature > 27: t, h = temperature, humidity heat_index = (-8.784 + 1.611*t + 2.339*h - 0.146*t*h - 0.012*t**2 - 0.016*h**2 + 0.002*t**2*h + 0.0007*t*h**2 - 3.6e-6*t**2*h**2) wind_chill = temperature if temperature < 10 and wind_speed > 4.8: wind_chill = (13.12 + 0.6215*temperature - 11.37*wind_speed**0.16 + 0.3965*temperature*wind_speed**0.16) storm_index = max(-pressure_trend, 0)*0.4 + wind_speed*0.3 + max(temperature-10, 0)*0.15 rain_potential = max((humidity-65)*0.8 + (cloud_cover-50)*0.3 + max(-pressure_trend,0)*2, 0) instability_idx = max(temperature-10, 0)*0.5 + max(humidity-60, 0)*0.3 cont_raw = np.array([ temperature, humidity, pressure, pressure_trend, wind_speed, altitude, cloud_cover, visibility, solar_radiation, dew_point, heat_index, wind_chill, storm_index, rain_potential, instability_idx, ], dtype=np.float32) m = np.array([MEANS[c] for c in CONTINUOUS_FEATURES], dtype=np.float32) s = np.array([STDS[c] for c in CONTINUOUS_FEATURES], dtype=np.float32) cont_norm = (cont_raw - m) / (s + 1e-8) inputs = { 'continuous': torch.tensor(cont_norm).unsqueeze(0), 'wind_direction': torch.tensor([WIND_DIR_MAP[wind_direction]]), 'month': torch.tensor([int(month)]), 'air_mass': torch.tensor([AIR_MASS_MAP[air_mass]]), 'climate_zone': torch.tensor([CLIMATE_ZONE_MAP[climate_zone]]), 'terrain_type': torch.tensor([TERRAIN_MAP[terrain_type]]), } with torch.no_grad(): probs = torch.softmax(model(**inputs)['logits'], dim=-1).squeeze(0).numpy() top5 = np.argsort(probs)[::-1][:5] pred = int(top5[0]) label = f"**{ID2LABEL[pred].replace('_',' ').title()}** ({probs[pred]*100:.1f}%)" derived = (f"Dew point: {dew_point:.1f}°C | Heat index: {heat_index:.1f}°C | " f"Wind chill: {wind_chill:.1f}°C | Storm index: {storm_index:.1f} | " f"Rain potential: {rain_potential:.1f}") bars = "" for idx in top5: pct = probs[idx] * 100 name = ID2LABEL[idx].replace('_', ' ').title() bars += (f'
' f'{name}' f'' f' {pct:.1f}%
') return label, bars, derived # ── Gradio UI ───────────────────────────────────────────────────────────────── with gr.Blocks(title="SupraWeather-Nano-1.5", theme=gr.themes.Soft()) as demo: gr.Markdown("# SupraWeather-Nano-1.5\nProcedural Atmospheric Reasoning · 25 Phenomena · FT-Transformer · SupraLabs") with gr.Row(): with gr.Column(): temp = gr.Slider(-50, 55, value=15, label="Temperature (°C)", step=0.5) hum = gr.Slider(0, 100, value=65, label="Humidity (%)", step=1) pres = gr.Slider(870,1060, value=1013, label="Pressure (hPa)", step=0.5) pt = gr.Slider(-20, 20, value=0, label="Pressure Trend (hPa/3h)",step=0.5) ws = gr.Slider(0, 120, value=10, label="Wind Speed (km/h)", step=1) alt = gr.Slider(0, 4500, value=100, label="Altitude (m)", step=50) cc = gr.Slider(0, 100, value=40, label="Cloud Cover (%)", step=1) vis = gr.Slider(0.05, 50, value=15, label="Visibility (km)", step=0.1) sol = gr.Slider(0, 1100, value=500, label="Solar Radiation (W/m²)", step=10) with gr.Column(): wd = gr.Dropdown(list(WIND_DIR_MAP.keys()), value='north', label="Wind Direction") mo = gr.Slider(1, 12, value=6, step=1, label="Month") am = gr.Dropdown(list(AIR_MASS_MAP.keys()), value='continental', label="Air Mass") cz = gr.Dropdown(list(CLIMATE_ZONE_MAP.keys()),value='temperate', label="Climate Zone") ter = gr.Dropdown(list(TERRAIN_MAP.keys()), value='flat', label="Terrain Type") gr.Markdown("---") pred_out = gr.Markdown("—") bars_out = gr.HTML() deriv_out = gr.Markdown() btn = gr.Button("Predict", variant="primary") btn.click(derive_and_predict, inputs=[temp,hum,pres,pt,ws,alt,cc,vis,sol,wd,mo,am,cz,ter], outputs=[pred_out, bars_out, deriv_out]) demo.launch()