Views
No views yet
| Pred 0 | Pred 1 | |
|---|---|---|
| True 0 | 0 | 2 |
| True 1 | 0 | 2 |
1import torch
2from transformers import AutoModel
3
4repo = "Wonder-Griffin/TorNet-Oracle"
5model = AutoModel.from_pretrained(repo, trust_remote_code=True).eval()
6
7# Example dummy batch
8B, T, H, W = 2, 1, 256, 256 # T time steps -> in_channels = 3*T (reflectivity, velocity, spectrum width?)
9radar_x = torch.randn(B, 3*T, H, W)
10
11# Atmospheric dictionary (use only what you have; shapes must be (B, dim))
12atmo = {
13 "cape": torch.randn(B, 1),
14 "wind_shear": torch.randn(B, 4), # 0–1, 0–3, 0–6, deep
15 "helicity": torch.randn(B, 2), # 0–1, 0–3
16 "temperature": torch.randn(B, 3), # sfc, 850, 500
17 "dewpoint": torch.randn(B, 2), # sfc, 850
18 "pressure": torch.randn(B, 1),
19}
20
21out = model(radar_x=radar_x, atmo=atmo)
22print(out.tornado_probability.shape) # (B,)
23print(out.ef_scale_probs.shape) # (B, 6)
24print(out.location_offset.shape) # (B, 2)
25print(out.timing_predictions.shape) # (B, 3)
26---
27
28# 3) Notes to avoid common gotchas
29
30- **Export the class names**: Make sure `StormOracleModel` and `StormOracleConfig` are importable at the repo root via `__init__.py`. Hugging Face uses that when `trust_remote_code=True`.
31- **Architectures**: The `"architectures"` array in `config.json` **must** include `"StormOracleModel"`.
32- **Weights**: You already have `pytorch_model.bin`/**or** `model.safetensors`. Either is fine. Keep the filenames standard.
33- **Forward signature**: With remote code, it’s okay that `forward` takes `radar_x` and `atmo`. Users pass them as keyword args as shown.
34- **Version pins**: If you rely on features from newer `transformers`, keep the `transformers_version` in `config.json` current.
35
36---
37
38# 4) Optional niceties
39
40- **`hubconf.py`** (for `torch.hub` users):
41 ```python
42 from .tornado_predictor import TornadoSuperPredictor
43
44 def storm_oracle(in_channels=3, pretrained=False, hf_repo=None, map_location="cpu"):
45 model = TornadoSuperPredictor(in_channels=in_channels)
46 if pretrained and hf_repo is not None:
47 from huggingface_hub import hf_hub_download
48 path = hf_hub_download(hf_repo, filename="pytorch_model.bin")
49 import torch
50 state = torch.load(path, map_location=map_location)
51 model.load_state_dict(state, strict=True)
52 return model