Update README.md

README.md

@@ -13,15 +13,18 @@ license: mit
 
 # Yield Estimation Transformer
 
-A Hugging Face Transformers model for county-level corn yield estimation using multi-temporal weather observations and static soil properties.
+A Hugging Face custom pipeline using Transformers model for county-level corn yield estimation using multi-temporal weather observations and static soil properties.
 
 The model combines weekly weather time-series with static soil features to estimate corn yield in bushels per acre (`bu/acre`). It is packaged for inference using Hugging Face Transformers and has been tested for deployment through FlexServ.
 
-The Hugging Face `text-classification` task is used as the FlexServ-compatible serving interface. The underlying model performs scalar regression, and the returned `score` represents predicted corn yield in `bu/acre`.
+The Hugging Face `text-classification` task is used as the FlexServ-compatible serving interface. The underlying model performs regression, and the returned `score` represents predicted corn yield in `bu/acre`.
+
+The `training_code` directory composes the source code for data preparation, model training, evaluation, and inference.
 
 ### Tags
+**Tags:** Digital-Agriculture, Foundation-AI
 
-- Crop Yield Estimation
+- In-season Crop Yield Estimation
 - Digital Agriculture
 - Transformers
 - Multi-Temporal Modeling
@@ -29,9 +32,12 @@ The Hugging Face `text-classification` task is used as the FlexServ-compatible s
 - Hugging Face Transformers
 - FlexServ
 
+For guidance on what to include in Tutorials, How-To Guides, Explanation, and Reference, see [Diátaxis](https://diataxis.fr/).
+
 ### License
 
-- [![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![License](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+This project is released under the MIT License. The full text is available in [LICENSE](LICENSE).
 
 ## References
 
@@ -64,7 +70,7 @@ FlexServ documentation: https://zhangwei217245.github.io/FlexServ/
 
 This work was developed as part of the ICICLE AI Institute.
 
-*National Science Foundation (NSF) AI Institute for Intelligent Cyberinfrastructure with Computational Learning in the Environment (ICICLE), Award OAC-2112606.*
+*National Science Foundation (NSF) funded AI institute for Intelligent Cyberinfrastructure with Computational Learning in the Environment (ICICLE) (OAC 2112606)*
 
 ## Issue reporting
 
@@ -80,7 +86,7 @@ Sarikaa Sridhar: sridhar.86@buckeyemail.osu.edu
 
 ### Overview
 
-The Yield Estimation Transformer is a pretrained model for county-level corn yield estimation. It combines multi-temporal weekly weather observations with static soil properties and produces a scalar yield prediction in bushels per acre.
+The Yield Estimation Transformer is a pretrained model for in-season county-level corn yield estimation. It combines temporal weekly weather observations with static soil properties and produces a scalar yield prediction in bushels per acre.
 
 The model accepts six weekly weather variables:
 
@@ -103,20 +109,6 @@ A cutoff determines how many weeks of weather information are available to the m
 
 For deployment through FlexServ, the model uses the Hugging Face `text-classification` pipeline as its serving interface. This is an interface choice for inference compatibility; the underlying prediction task remains regression.
 
-### Prerequisites
-
-- Python 3.10+
-- PyTorch
-- Hugging Face Transformers
-- Dependencies listed in `requirements.txt`
-- FlexServ environment for service deployment
-
-Because the repository provides custom model configuration, tokenizer, and architecture code, Hugging Face loading requires:
-
-```python
-trust_remote_code=True
-```
-
 ---
 
 # How-To Guides
@@ -154,6 +146,9 @@ The repository contains the files required for standalone Hugging Face and FlexS
 ```text
 .
 ├── README.md
+├── LICENSE
+├── component-info.yaml
+├── training_code/
 ├── config.json
 ├── configuration_yield.py
 ├── model.safetensors
@@ -171,26 +166,10 @@ A complete inference example is provided in:
 sample_input_weekly.json
 ```
 
-### Installation
-
-Clone the model repository:
-
-```bash
-git clone https://huggingface.co/ICICLE-AI/yield-estimation
-cd yield-estimation
-```
-
-Create and activate a Python environment:
-
-```bash
-conda create -n yield_hf python=3.10
-conda activate yield_hf
-```
-
-Install the required dependencies:
+Because the repository provides custom model configuration, tokenizer, and architecture code, Hugging Face loading requires:
 
-```bash
-pip install -r requirements.txt
+```python
+trust_remote_code=True
 ```
 
 ### Usage
@@ -329,6 +308,352 @@ Expected output for the included sample is approximately:
 [{'label': 'YIELD_BU_ACRE', 'score': 165.1769561767578}]
 ```
 
+### Training own model
+
+#### Installation
+
+Clone the model repository:
+
+```bash
+git clone https://huggingface.co/ICICLE-AI/yield-estimation
+cd yield-estimation/training_code
+```
+
+Create and activate a Python environment:
+
+```bash
+conda create -n yield_hf python=3.10
+conda activate yield_hf
+```
+
+Install the required dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+#### Prerequisites
+
+- Python 3.10+
+- PyTorch
+- Hugging Face Transformers
+- Dependencies listed in `requirements.txt`
+
+For GPU training, a CUDA-compatible PyTorch environment is recommended.
+
+The included SLURM script can be used on a compatible HPC system.
+
+#### Problem Description
+
+The objective is to estimate county-level corn yield from weather and soil information.
+
+Each sample contains two primary sources of information:
+
+1. multi-temporal weather observations;
+2. static soil properties.
+
+The weather input is represented as:
+
+```text
+[K, W]
+```
+
+where:
+
+- `K` is the number of temporal observations;
+- `W = 6` is the number of weather variables.
+
+The static soil input is represented as:
+
+```text
+[S]
+```
+
+where:
+
+```text
+S = 66
+```
+
+The model learns a regression function of the form:
+
+```text
+weather + soil + crop information
+                ↓
+        transformer model
+                ↓
+       predicted corn yield
+```
+
+The predicted value is expressed in bushels per acre (`bu/acre`).
+
+The model supports multi-temporal estimation by evaluating the weather sequence at different seasonal cutoffs.
+
+#### Repository Structure
+
+The final repository is organized as follows:
+
+```text
+.
+├── README.md
+├── requirements.txt
+├── training.slurm
+│
+├── checkpoints/
+│   ├── config.json
+│   ├── metrics.json
+│   └── model.safetensors
+│
+├── config/
+│   ├── __init__.py
+│   └── config.py
+│
+├── data/
+│   ├── __init__.py
+│   ├── dataset.py
+│   └── preprocessing.py
+│
+├── examples/
+│   └── sample_input_weekly.json
+│
+├── hf/
+│   ├── __init__.py
+│   ├── auto.py
+│   ├── configuration_yield.py
+│   └── modeling_yield.py
+│
+├── models/
+│   ├── __init__.py
+│   └── unimodal_ws_crossattn.py
+│
+├── scripts/
+│   ├── __init__.py
+│   ├── prepare_cornbelt.py
+│   ├── train_hf.py
+│   ├── evaluate_hf.py
+│   └── inference_hf.py
+│
+└── training/
+    ├── __init__.py
+    └── engine.py
+```
+
+The major components are:
+
+- `data/` — dataset loading and preprocessing
+- `models/` — core neural network architecture
+- `training/` — training and evaluation utilities
+- `hf/` — Hugging Face AutoClass-compatible regression wrapper used by the training repository
+- `scripts/` — data preparation, training, evaluation, and inference entry points
+- `checkpoints/` — final trained checkpoint and configuration
+- `examples/` — example structured model input
+- `training.slurm` — example HPC training job
+
+#### Data Preparation
+
+The USA County Level Crop Yield public dataset is used for training, validation and testing. The data preparation workflow is implemented in:
+
+```text
+scripts/prepare_cornbelt.py
+```
+
+After preparation, the expected dataset structure is:
+
+```text
+data/
+└── cornbelt/
+    ├── train.h5
+    ├── val.h5
+    └── test.h5
+```
+
+The model uses the following six weather variables:
+
+```text
+prcp
+srad
+swe
+tmax
+tmin
+vp
+```
+
+The 66 soil variables used by the final checkpoint are recorded in the model configuration.
+
+#### Training
+
+The primary training entry point is:
+
+```text
+scripts/train_hf.py
+```
+
+The final model uses multi-cutoff training with:
+
+```text
+20,24,28,32,36,40,44,48,52
+```
+
+An example training command is:
+
+```bash
+python scripts/train_hf.py \
+  --train_file data/cornbelt/train.h5 \
+  --val_file data/cornbelt/val.h5 \
+  --test_file data/cornbelt/test.h5 \
+  --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 \
+  --crop corn \
+  --time_agg weekly \
+  --train_cutoffs 20,24,28,32,36,40,44,48,52 \
+  --eval_cutoffs 20,24,28,32,36,40,44,48,52 \
+  --epochs 30 \
+  --lr 3e-5 \
+  --batch_size 32 \
+  --out_dir checkpoints
+```
+
+The final checkpoint is stored in:
+
+```text
+checkpoints/
+```
+
+The checkpoint includes:
+
+```text
+config.json
+model.safetensors
+metrics.json
+```
+#### Training on Your Own Data
+
+The training pipeline can also be used to train a new yield estimation model on a compatible dataset.
+
+Prepare the dataset in the HDF5 format expected by `YieldDataset` and provide separate training, validation, and test files.
+
+The weather and soil variables supplied to the training command must correspond to the variables available in the prepared dataset.
+
+A general training command is:
+
+```bash
+python scripts/train_hf.py \
+  --train_file <path/to/train.h5> \
+  --val_file <path/to/val.h5> \
+  --test_file <path/to/test.h5> \
+  --weather_vars <comma-separated-weather-variables> \
+  --soil_vars <comma-separated-soil-variables> \
+  --crop <crop-name> \
+  --time_agg weekly \
+  --train_cutoffs <comma-separated-training-cutoffs> \
+  --eval_cutoffs <comma-separated-evaluation-cutoffs> \
+  --epochs <number-of-epochs> \
+  --lr <learning-rate> \
+  --batch_size <batch-size> \
+  --out_dir <output-directory>
+```
+
+#### SLURM Training
+
+An example SLURM job is provided in:
+
+```text
+training.slurm
+```
+
+Submit it using:
+
+```bash
+sbatch training.slurm
+```
+
+#### Evaluation
+
+The trained checkpoint can be evaluated using:
+
+```text
+scripts/evaluate_hf.py
+```
+
+For the final multi-cutoff model:
+
+```bash
+python scripts/evaluate_hf.py \
+  --hf_model_dir checkpoints \
+  --test_file data/cornbelt/test.h5 \
+  --cutoffs 20,24,28,32,36,40,44,48,52 \
+  --batch_size 64 \
+  --output_csv checkpoints/test_predictions.csv \
+  --metrics_json checkpoints/test_metrics.json
+```
+
+Evaluation is performed independently at the configured seasonal cutoffs.
+
+The evaluation process:
+
+1. loads the final trained checkpoint;
+2. loads the test dataset;
+3. applies the normalization statistics stored in the checkpoint configuration;
+4. performs inference at the requested cutoffs;
+5. computes evaluation metrics;
+6. save predictions and metrics to disk.
+
+#### Inference
+
+Inference using the trained checkpoint is implemented in:
+
+```text
+scripts/inference_hf.py
+```
+
+An example structured input is provided in:
+
+```text
+examples/sample_input_weekly.json
+```
+
+The sample follows the general structure:
+
+```json
+{
+  "crop": "corn",
+  "weather_format": "weekly",
+  "cutoff": 52,
+  "weather": {
+    "prcp": [],
+    "srad": [],
+    "swe": [],
+    "tmax": [],
+    "tmin": [],
+    "vp": []
+  },
+  "soil": {
+    "bdod_mean_0-5cm": 0.0
+  }
+}
+```
+
+The complete sample file contains the required weather sequence and soil variables.
+
+Run single-sample inference with:
+
+```bash
+python scripts/inference_hf.py \
+  --hf_model_dir checkpoints \
+  --single_sample_json examples/sample_input_weekly.json \
+  --cutoff 52 \
+  --output_csv inference_prediction.csv
+```
+
+The output contains the predicted yield for the requested cutoff.
+
+For example:
+
+```text
+sample_idx,cutoff,y_pred
+0,52,<predicted_yield>
+```
+
 ---
 
 # Explanation