inference.py

import joblib
import numpy as np
import pandas as pd

# Load all saved artifacts
model = joblib.load("rf_yield.pkl")
scaler = joblib.load("scaler.pkl")
ohe = joblib.load("ohe.pkl")
training_columns = joblib.load("training_columns.pkl")

def predict(inputs: dict):
    """
    inputs example:
    {"area": 50, "season": "Kharif"}
    """

    # Convert input to dataframe
    df = pd.DataFrame([inputs])

    # Handle categorical encoding
    categorical_cols = ["season"]   # change/add depending on your dataset
    encoded = ohe.transform(df[categorical_cols]).toarray()
    encoded_df = pd.DataFrame(encoded, columns=ohe.get_feature_names_out(categorical_cols))

    # Keep numerical features
    numeric_df = df.drop(columns=categorical_cols)

    # Combine
    full_df = pd.concat([numeric_df, encoded_df], axis=1)

    # Reorder columns to match training
    full_df = full_df.reindex(columns=training_columns, fill_value=0)

    # Scale
    scaled = scaler.transform(full_df)

    # Predict
    pred = model.predict(scaled)

    return {"prediction": str(pred[0])}