| 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" |
|
|
| |
| 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} |
|
|
| |
| 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} |
|
|
| |
| 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 |
|
|
| |
| 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'<div style="margin:6px 0;font-family:monospace">' |
| f'<span style="display:inline-block;width:160px">{name}</span>' |
| f'<span style="display:inline-block;width:{int(pct*3)}px;height:14px;' |
| f'background:#3498db;vertical-align:middle;border-radius:3px"></span>' |
| f' {pct:.1f}%</div>') |
| return label, bars, derived |
|
|
| |
| 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() |
