Upload 4 files

config.json

@@ -3,7 +3,7 @@
   "S": 66,
   "W": 6,
   "architectures": [
-    "YieldForRegression"
+    "YieldForSequenceClassification"
   ],
   "crop_emb_dim": 8,
   "d_model": 64,
@@ -272,8 +272,22 @@
   "auto_map": {
     "AutoConfig": "configuration_yield.YieldConfig",
     "AutoModel": "modeling_yield.YieldForRegression",
+    "AutoModelForSequenceClassification": "modeling_yield_text.YieldForSequenceClassification",
+    "AutoTokenizer": [
+      "tokenization_yield.YieldTokenizer",
+      null
+    ],
     "AutoProcessor": "processing_yield.YieldProcessor"
   },
+  "problem_type": "regression",
+  "num_labels": 1,
+  "id2label": {
+    "0": "YIELD_BU_ACRE"
+  },
+  "label2id": {
+    "YIELD_BU_ACRE": 0
+  },
+  "tokenizer_class": "YieldTokenizer",
   "custom_pipelines": {
     "yield-estimation": {
       "impl": "pipeline_yield.YieldEstimationPipeline",

modeling_yield_text.py

@@ -0,0 +1,209 @@
+import torch
+
+from torch import nn
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+from .configuration_yield import YieldConfig
+from .yield_transformer import YieldTransformer
+
+
+class YieldForSequenceClassification(PreTrainedModel):
+
+    config_class = YieldConfig
+    base_model_prefix = "yield_model"
+
+    def __init__(self, config: YieldConfig):
+        super().__init__(config)
+
+        self.yield_model = YieldTransformer(
+            w_dim=config.W,
+            soil_dim=config.S,
+            d_model=config.d_model,
+            nhead=config.nhead,
+            num_layers=config.num_layers,
+            dim_ff=config.dim_ff,
+            dropout=config.dropout,
+            use_crop=config.use_crop,
+            crop_emb_dim=config.crop_emb_dim,
+            max_weeks=max(52, config.K),
+            pool=config.pool,
+        )
+
+        self.post_init()
+
+    def forward(
+        self,
+        weather,
+        soil,
+        crop_id,
+        horizon_idx=None,
+        labels=None,
+        **kwargs,
+    ):
+
+        # ==================================================
+        # Shape checks
+        # ==================================================
+
+        if weather.ndim != 3:
+            raise ValueError(
+                f"weather must have shape [B,K,W], "
+                f"received {tuple(weather.shape)}"
+            )
+
+        if soil.ndim != 2:
+            raise ValueError(
+                f"soil must have shape [B,S], "
+                f"received {tuple(soil.shape)}"
+            )
+
+        # ==================================================
+        # Training normalization statistics
+        # ==================================================
+
+        w_mean = torch.tensor(
+            self.config.w_mean,
+            device=weather.device,
+            dtype=weather.dtype,
+        )
+
+        w_std = torch.tensor(
+            self.config.w_std,
+            device=weather.device,
+            dtype=weather.dtype,
+        )
+
+        s_mean = torch.tensor(
+            self.config.s_mean,
+            device=soil.device,
+            dtype=soil.dtype,
+        )
+
+        s_std = torch.tensor(
+            self.config.s_std,
+            device=soil.device,
+            dtype=soil.dtype,
+        )
+
+        # NaN -> training mean
+        weather = torch.where(
+            torch.isnan(weather),
+            w_mean.view(1, 1, -1),
+            weather,
+        )
+
+        soil = torch.where(
+            torch.isnan(soil),
+            s_mean.view(1, -1),
+            soil,
+        )
+
+        # normalize
+        weather = (
+            weather - w_mean.view(1, 1, -1)
+        ) / w_std.view(1, 1, -1)
+
+        soil = (
+            soil - s_mean.view(1, -1)
+        ) / s_std.view(1, -1)
+
+        # ==================================================
+        # Cutoff
+        # ==================================================
+
+        if horizon_idx is None:
+            horizon_idx = torch.full(
+                (weather.shape[0],),
+                weather.shape[1],
+                dtype=torch.long,
+                device=weather.device,
+            )
+
+        if not torch.is_tensor(horizon_idx):
+            horizon_idx = torch.tensor(
+                horizon_idx,
+                dtype=torch.long,
+                device=weather.device,
+            )
+
+        horizon_idx = horizon_idx.to(
+            weather.device
+        ).long()
+
+        if horizon_idx.ndim == 0:
+            horizon_idx = horizon_idx.unsqueeze(0)
+
+        # FlexServe requests are normally one sample.
+        # Ensure one common temporal length for a batch.
+        unique_cutoffs = torch.unique(
+            horizon_idx
+        )
+
+        if len(unique_cutoffs) != 1:
+            raise ValueError(
+                "All samples in one batch must use the same cutoff."
+            )
+
+        t_eff = int(
+            unique_cutoffs[0].item()
+        )
+
+        weather = weather[
+            :,
+            :t_eff,
+            :
+        ]
+
+        # ==================================================
+        # Existing trained model
+        # ==================================================
+
+        logits_norm = self.yield_model(
+            weather,
+            soil,
+            crop_id,
+            horizon_idx=horizon_idx,
+            causal=True,
+            return_sequence=False,
+        )
+
+        # Restore yield to bu/acre.
+        y_mean = torch.tensor(
+            self.config.y_mean,
+            device=logits_norm.device,
+            dtype=logits_norm.dtype,
+        )
+
+        y_std = torch.tensor(
+            self.config.y_std,
+            device=logits_norm.device,
+            dtype=logits_norm.dtype,
+        )
+
+        predicted_yield = (
+            logits_norm * y_std
+            + y_mean
+        )
+
+        # TextClassificationPipeline expects [B, num_labels].
+        logits = predicted_yield.unsqueeze(-1)
+
+        loss = None
+
+        if labels is not None:
+
+            labels = labels.to(
+                logits.dtype
+            ).view(-1)
+
+            loss = nn.functional.mse_loss(
+                predicted_yield,
+                labels,
+            )
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+        )

tokenization_yield.py

@@ -0,0 +1,291 @@
+import json
+
+import numpy as np
+import torch
+
+from transformers import PreTrainedTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+
+
+class YieldTokenizer(PreTrainedTokenizer):
+    """
+    Adapter tokenizer for FlexServe's built-in text-classification pipeline.
+
+    This does NOT tokenize natural language.
+
+    It accepts:
+      - a JSON string, or
+      - the yield input dictionary
+
+    and converts it into:
+      weather:     [B, 52, 6]
+      soil:        [B, 66]
+      crop_id:     [B]
+      horizon_idx: [B]
+
+    Normalization is intentionally NOT performed here.
+    The sequence-classification model wrapper performs normalization
+    using statistics stored in config.json.
+    """
+
+    vocab_files_names = {}
+    model_input_names = [
+        "weather",
+        "soil",
+        "crop_id",
+        "horizon_idx",
+    ]
+
+    def __init__(
+        self,
+        weather_vars=None,
+        soil_vars=None,
+        K=52,
+        eval_cutoffs=None,
+        **kwargs,
+    ):
+        self.weather_vars = list(weather_vars or [])
+        self.soil_vars = list(soil_vars or [])
+        self.K = int(K)
+        self.eval_cutoffs = list(
+            eval_cutoffs
+            or [20, 24, 28, 32, 36, 40, 44, 48, 52]
+        )
+
+        super().__init__(
+            pad_token="[PAD]",
+            unk_token="[UNK]",
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self):
+        return 2
+
+    def get_vocab(self):
+        return {
+            "[PAD]": 0,
+            "[UNK]": 1,
+        }
+
+    def _tokenize(self, text, **kwargs):
+        return ["[UNK]"]
+
+    def _convert_token_to_id(self, token):
+        return 0 if token == "[PAD]" else 1
+
+    def _convert_id_to_token(self, index):
+        return "[PAD]" if index == 0 else "[UNK]"
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        return ()
+
+    def _parse_sample(self, sample):
+
+        if isinstance(sample, str):
+            sample = sample.strip()
+
+            try:
+                sample = json.loads(sample)
+            except json.JSONDecodeError as exc:
+                raise ValueError(
+                    "Input must be a valid JSON string."
+                ) from exc
+
+            # Handle a JSON string containing another JSON string.
+            if isinstance(sample, str):
+                try:
+                    sample = json.loads(sample)
+                except json.JSONDecodeError as exc:
+                    raise ValueError(
+                        "Input string does not contain valid yield JSON."
+                    ) from exc
+
+        if not isinstance(sample, dict):
+            raise ValueError(
+                "Yield input must be a JSON object/dictionary."
+            )
+
+        crop = str(
+            sample.get("crop", "corn")
+        ).strip().lower()
+
+        if crop not in ("corn", "maize"):
+            raise ValueError(
+                "This released model supports corn only."
+            )
+
+        if "weather" not in sample:
+            raise ValueError(
+                "Missing 'weather' object."
+            )
+
+        if "soil" not in sample:
+            raise ValueError(
+                "Missing 'soil' object."
+            )
+
+        weather_dict = sample["weather"]
+        soil_dict = sample["soil"]
+
+        # --------------------------------------------
+        # Weather: [52, 6]
+        # --------------------------------------------
+
+        weather_cols = []
+
+        for var in self.weather_vars:
+
+            if var not in weather_dict:
+                raise ValueError(
+                    f"Missing weather variable '{var}'."
+                )
+
+            values = np.asarray(
+                weather_dict[var],
+                dtype=np.float32,
+            )
+
+            if values.ndim != 1:
+                raise ValueError(
+                    f"Weather '{var}' must be one-dimensional."
+                )
+
+            if len(values) != self.K:
+                raise ValueError(
+                    f"Weather '{var}' requires exactly "
+                    f"{self.K} weekly values; received {len(values)}."
+                )
+
+            weather_cols.append(values)
+
+        weather = np.stack(
+            weather_cols,
+            axis=1,
+        ).astype(np.float32)
+
+        # --------------------------------------------
+        # Soil: [66]
+        # --------------------------------------------
+
+        soil = []
+
+        for var in self.soil_vars:
+
+            if var not in soil_dict:
+                raise ValueError(
+                    f"Missing soil variable '{var}'."
+                )
+
+            soil.append(
+                float(soil_dict[var])
+            )
+
+        soil = np.asarray(
+            soil,
+            dtype=np.float32,
+        )
+
+        # --------------------------------------------
+        # Cutoff
+        # --------------------------------------------
+
+        cutoff = int(
+            sample.get(
+                "cutoff",
+                max(self.eval_cutoffs),
+            )
+        )
+
+        if cutoff not in self.eval_cutoffs:
+            raise ValueError(
+                f"Unsupported cutoff {cutoff}. "
+                f"Supported cutoffs are {self.eval_cutoffs}."
+            )
+
+        return {
+            "weather": weather,
+            "soil": soil,
+            "crop_id": 0,
+            "horizon_idx": cutoff,
+        }
+
+    def __call__(
+        self,
+        text=None,
+        text_pair=None,
+        return_tensors=None,
+        **kwargs,
+    ):
+        # --------------------------------------------------
+        # FlexServe/HF may call tokenizer(**input_dict)
+        # instead of tokenizer(json_string).
+        # --------------------------------------------------
+
+        if text is None and "weather" in kwargs and "soil" in kwargs:
+
+            sample = {
+                "crop": kwargs.pop("crop", "corn"),
+                "weather": kwargs.pop("weather"),
+                "soil": kwargs.pop("soil"),
+                "cutoff": kwargs.pop(
+                    "cutoff",
+                    max(self.eval_cutoffs),
+                ),
+            }
+
+            samples = [sample]
+
+        elif isinstance(text, (list, tuple)):
+
+            samples = list(text)
+
+        else:
+
+            samples = [text]
+
+        parsed = [
+            self._parse_sample(sample)
+            for sample in samples
+        ]
+
+        weather = torch.tensor(
+            np.stack(
+                [x["weather"] for x in parsed],
+                axis=0,
+            ),
+            dtype=torch.float32,
+        )
+
+        soil = torch.tensor(
+            np.stack(
+                [x["soil"] for x in parsed],
+                axis=0,
+            ),
+            dtype=torch.float32,
+        )
+
+        crop_id = torch.tensor(
+            [
+                x["crop_id"]
+                for x in parsed
+            ],
+            dtype=torch.long,
+        )
+
+        horizon_idx = torch.tensor(
+            [
+                x["horizon_idx"]
+                for x in parsed
+            ],
+            dtype=torch.long,
+        )
+
+        return BatchEncoding(
+            {
+                "weather": weather,
+                "soil": soil,
+                "crop_id": crop_id,
+                "horizon_idx": horizon_idx,
+            }
+        )

tokenizer_config.json

@@ -0,0 +1,98 @@
+{
+  "tokenizer_class": "YieldTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_yield.YieldTokenizer",
+      null
+    ]
+  },
+  "weather_vars": [
+    "prcp",
+    "srad",
+    "swe",
+    "tmax",
+    "tmin",
+    "vp"
+  ],
+  "soil_vars": [
+    "bdod_mean_0-5cm",
+    "bdod_mean_5-15cm",
+    "bdod_mean_15-30cm",
+    "bdod_mean_30-60cm",
+    "bdod_mean_60-100cm",
+    "bdod_mean_100-200cm",
+    "cec_mean_0-5cm",
+    "cec_mean_5-15cm",
+    "cec_mean_15-30cm",
+    "cec_mean_30-60cm",
+    "cec_mean_60-100cm",
+    "cec_mean_100-200cm",
+    "cfvo_mean_0-5cm",
+    "cfvo_mean_5-15cm",
+    "cfvo_mean_15-30cm",
+    "cfvo_mean_30-60cm",
+    "cfvo_mean_60-100cm",
+    "cfvo_mean_100-200cm",
+    "clay_mean_0-5cm",
+    "clay_mean_5-15cm",
+    "clay_mean_15-30cm",
+    "clay_mean_30-60cm",
+    "clay_mean_60-100cm",
+    "clay_mean_100-200cm",
+    "nitrogen_mean_0-5cm",
+    "nitrogen_mean_5-15cm",
+    "nitrogen_mean_15-30cm",
+    "nitrogen_mean_30-60cm",
+    "nitrogen_mean_60-100cm",
+    "nitrogen_mean_100-200cm",
+    "ocd_mean_0-5cm",
+    "ocd_mean_5-15cm",
+    "ocd_mean_15-30cm",
+    "ocd_mean_30-60cm",
+    "ocd_mean_60-100cm",
+    "ocd_mean_100-200cm",
+    "ocs_mean_0-5cm",
+    "ocs_mean_5-15cm",
+    "ocs_mean_15-30cm",
+    "ocs_mean_30-60cm",
+    "ocs_mean_60-100cm",
+    "ocs_mean_100-200cm",
+    "phh2o_mean_0-5cm",
+    "phh2o_mean_5-15cm",
+    "phh2o_mean_15-30cm",
+    "phh2o_mean_30-60cm",
+    "phh2o_mean_60-100cm",
+    "phh2o_mean_100-200cm",
+    "sand_mean_0-5cm",
+    "sand_mean_5-15cm",
+    "sand_mean_15-30cm",
+    "sand_mean_30-60cm",
+    "sand_mean_60-100cm",
+    "sand_mean_100-200cm",
+    "silt_mean_0-5cm",
+    "silt_mean_5-15cm",
+    "silt_mean_15-30cm",
+    "silt_mean_30-60cm",
+    "silt_mean_60-100cm",
+    "silt_mean_100-200cm",
+    "soc_mean_0-5cm",
+    "soc_mean_5-15cm",
+    "soc_mean_15-30cm",
+    "soc_mean_30-60cm",
+    "soc_mean_60-100cm",
+    "soc_mean_100-200cm"
+  ],
+  "K": 52,
+  "eval_cutoffs": [
+    20,
+    24,
+    28,
+    32,
+    36,
+    40,
+    44,
+    48,
+    52
+  ],
+  "model_max_length": 1000000
+}