97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
import sys
|
|
import os
|
|
import logging
|
|
from training.storage_client import MinIOStorageClient
|
|
import pandas as pd
|
|
from sklearn.model_selection import train_test_split
|
|
from xgboost import XGBClassifier
|
|
import mlflow
|
|
import mlflow.sklearn
|
|
import joblib
|
|
|
|
# Setup logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def train():
|
|
# 1. Setup Environment
|
|
# MLflow server URL (Internal K8s DNS)
|
|
mlflow.set_tracking_uri("http://mlflow.geocrop.svc.cluster.local:5000")
|
|
mlflow.set_experiment("GeoCrop_Retraining")
|
|
|
|
storage = MinIOStorageClient()
|
|
|
|
# 2. Fetch Data
|
|
logger.info("Fetching latest training data...")
|
|
# For now, let's load batch_1.csv as an example
|
|
df = storage.load_dataset('geocrop-datasets', 'batch_1.csv')
|
|
|
|
if df is None or df.empty:
|
|
logger.error("No data found!")
|
|
return
|
|
|
|
# 3. Simple Feature Engineering (Placeholder)
|
|
# Filter columns to only include features and target
|
|
# FEATURE_ORDER_V1 (subset for example)
|
|
features = ['ndvi_peak', 'evi_peak', 'savi_peak']
|
|
target = 'class' # Update with actual column name
|
|
|
|
X = df[features]
|
|
y = df[target]
|
|
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
|
|
|
|
# 4. Training with MLflow Tracking
|
|
with mlflow.start_run():
|
|
logger.info("Training model...")
|
|
model = XGBClassifier(n_estimators=100, max_depth=5)
|
|
model.fit(X_train, y_train)
|
|
|
|
# Log Metrics
|
|
score = model.score(X_test, y_test)
|
|
mlflow.log_metric("accuracy", score)
|
|
logger.info(f"Accuracy: {score}")
|
|
|
|
# Log Model to MLflow
|
|
mlflow.sklearn.log_model(model, "model")
|
|
|
|
# 5. Save Artifact and Upload to MinIO
|
|
model_filename = "Zimbabwe_Ensemble_Model_latest.pkl"
|
|
joblib.dump(model, model_filename)
|
|
|
|
storage.upload_file(model_filename, 'geocrop-models', model_filename)
|
|
logger.info(f"Uploaded model to MinIO: {model_filename}")
|
|
|
|
# 6. Generate Inference Script (Dynamic)
|
|
generate_inference_script(features)
|
|
|
|
def generate_inference_script(features):
|
|
"""
|
|
Creates a new inference script tailored to the trained model's features.
|
|
"""
|
|
script_content = f"""
|
|
import os
|
|
from training.storage_client import MinIOStorageClient
|
|
import joblib
|
|
import pandas as pd
|
|
|
|
def run_inference(aoi_data):
|
|
storage = MinIOStorageClient()
|
|
# Download the latest model
|
|
storage.download_file('geocrop-models', 'Zimbabwe_Ensemble_Model_latest.pkl', 'model.pkl')
|
|
model = joblib.load('model.pkl')
|
|
|
|
# Preprocess aoi_data (should pull from DEA STAC here)
|
|
# features = {features}
|
|
# ... logic to fetch from DEA STAC and compute these features ...
|
|
|
|
# result = model.predict(preprocessed_data)
|
|
# return result
|
|
"""
|
|
with open("training/generated_inference.py", "w") as f:
|
|
f.write(script_content)
|
|
logger.info("Generated new inference script: training/generated_inference.py")
|
|
|
|
if __name__ == "__main__":
|
|
train()
|