inference.py

import joblib
import json
import numpy as np

# ========================
# LOAD ARTIFACTS
# ========================
lgb_model = joblib.load("lgb_model.pkl")
scaler = joblib.load("scaler.pkl")
le = joblib.load("label_encoder.pkl")

with open("feature_names.json") as f:
    feature_names = json.load(f)

# ========================
# MITRE MAPPING
# ========================
mitre_mapping = {
    "Web Attack Sql Injection": ("T1190", "Exploit Public-Facing Application"),
    "DDoS": ("T1498", "Network Denial of Service"),
    "PortScan": ("T1046", "Network Service Scanning"),
    "Brute Force": ("T1110", "Brute Force")
}

# ========================
# HF ENTRY POINT
# ========================
def predict(inputs):
    """
    Hugging Face expects:
    {"inputs": [...]}
    """

    input_data = inputs

    # Convert to numpy
    x = np.array(input_data).reshape(1, -1)

    # Scale
    x_scaled = scaler.transform(x)

    # Prediction
    pred = lgb_model.predict(x_scaled)[0]
    pred_label = le.inverse_transform([pred])[0]

    # Confidence
    proba = lgb_model.predict_proba(x_scaled)[0]
    confidence = float(np.max(proba))
    confidence = min(confidence, 0.99)

    # MITRE
    mitre_id, mitre_name = mitre_mapping.get(
        pred_label, ("Unknown", "Unknown")
    )

    return {
        "prediction": pred_label,
        "confidence": round(confidence, 3),
        "mitre_attack": {
            "technique_id": mitre_id,
            "technique_name": mitre_name
        }
    }