""" Stacking classifier inference for SBAN multi-representation malware dataset classification. Loads sban_weighted_stacking_model.joblib and predicts which source dataset (bodmas, dike, malwarebazaar, sorel20m) each sample belongs to. """ from __future__ import annotations import argparse import math import re from collections import Counter from pathlib import Path import joblib import numpy as np import pandas as pd from scipy.sparse import csr_matrix, hstack from sklearn.metrics import ( accuracy_score, classification_report, f1_score, ) PROJECT_DIR = Path(__file__).resolve().parent DEFAULT_MODEL_PATH = PROJECT_DIR / "sban_weighted_stacking_model.joblib" ID_COLUMN = "ID" LABEL_COLUMN = "dataset_name" REPRESENTATION_COLUMNS = { "asm": "assembly_code", "binary": "binary_code", "source": "source_code", "nld": "NLD", } NUMERIC_FEATURE_NAMES = [ "char_count", "line_count", "token_count", "avg_line_length", "unique_token_ratio", "char_entropy", ] def binary_to_byte_tokens(text: str) -> str: tokens: list[str] = [] for line in text.splitlines(): line = re.sub(r"\s+", "", line) if not line or re.fullmatch(r"[0-9a-fA-F]+", line) is None: continue byte_count = len(line) // 2 tokens.extend(line[i * 2 : (i + 1) * 2].lower() for i in range(byte_count)) tokens.append("instsep") return " ".join(tokens) def calculate_char_entropy(text: str) -> float: if not text: return 0.0 counts = Counter(text) length = len(text) return -sum( (count / length) * math.log2(count / length) for count in counts.values() ) def extract_numeric_features(text: str) -> list[float]: lines = text.splitlines() tokens = text.split() char_count = len(text) line_count = len(lines) token_count = len(tokens) avg_line_length = ( sum(len(line) for line in lines) / line_count if line_count else 0.0 ) unique_token_ratio = ( len(set(tokens)) / token_count if token_count else 0.0 ) return [ float(char_count), float(line_count), float(token_count), float(avg_line_length), float(unique_token_ratio), calculate_char_entropy(text), ] def build_numeric_features(series: pd.Series) -> np.ndarray: return np.asarray( [extract_numeric_features(text) for text in series.astype(str)], dtype=np.float32, ) def build_inference_features( representation_name: str, text_series: pd.Series, representation_bundle: dict, ) -> tuple[csr_matrix, np.ndarray]: texts = text_series.fillna("").astype(str).tolist() if representation_name == "binary": tfidf_texts = [binary_to_byte_tokens(text) for text in texts] else: tfidf_texts = texts vectorizer = representation_bundle["vectorizer"] numeric_scaler = representation_bundle["numeric_scaler"] x_tfidf = vectorizer.transform(tfidf_texts) x_numeric = build_numeric_features(text_series) x_numeric_scaled = numeric_scaler.transform(x_numeric) x_full = hstack( [x_tfidf, csr_matrix(x_numeric_scaled)], format="csr", dtype=np.float32, ) return x_full, np.asarray(x_numeric_scaled, dtype=np.float32) class StackingPredictor: """Wrapper around the exported joblib inference bundle.""" def __init__(self, model_path: Path | str = DEFAULT_MODEL_PATH) -> None: self.model_path = Path(model_path) self.model_bundle = joblib.load(self.model_path) self.representation_order = self.model_bundle["representation_order"] self.representation_columns = self.model_bundle["representation_columns"] self.label_encoder = self.model_bundle["label_encoder"] self.meta_model = self.model_bundle["meta_model"] self.representation_bundles = self.model_bundle["representations"] @property def classes(self) -> list[str]: return self.label_encoder.classes_.tolist() def predict(self, df: pd.DataFrame) -> pd.DataFrame: missing = [ col for rep in self.representation_order if (col := self.representation_columns[rep]) not in df.columns ] if missing: raise ValueError(f"Missing required columns: {missing}") meta_features: list[np.ndarray] = [] for representation_name in self.representation_order: bundle = self.representation_bundles[representation_name] x_full, x_numeric = build_inference_features( representation_name, df[self.representation_columns[representation_name]], bundle, ) x_selected = x_full[:, bundle["selected_indices"]] decision_scores = bundle["base_model"].decision_function(x_selected) meta_features.append(decision_scores) meta_features.append(x_numeric) x_meta = np.hstack(meta_features) encoded_predictions = self.meta_model.predict(x_meta) prediction_probabilities = self.meta_model.predict_proba(x_meta) result = df.copy() result["prediction"] = self.label_encoder.inverse_transform( encoded_predictions ) for index, class_name in enumerate(self.label_encoder.classes_): result[f"prob_{class_name}"] = prediction_probabilities[:, index] return result def load_predictor(model_path: Path | str = DEFAULT_MODEL_PATH) -> StackingPredictor: return StackingPredictor(model_path) def evaluate_predictions( predictions: pd.DataFrame, label_column: str = LABEL_COLUMN, ) -> None: if label_column not in predictions.columns: print(f"Skip evaluation: column {label_column!r} not in input.") return y_true = predictions[label_column] y_pred = predictions["prediction"] print(f"Accuracy : {accuracy_score(y_true, y_pred):.6f}") print(f"Macro F1 : {f1_score(y_true, y_pred, average='macro'):.6f}") print(f"Weighted F1 : {f1_score(y_true, y_pred, average='weighted'):.6f}") print("\nClassification Report\n") print( classification_report( y_true, y_pred, digits=4, target_names=sorted(y_true.unique()), ) ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Run SBAN weighted stacking inference on a parquet file.", ) parser.add_argument( "--model-path", type=Path, default=DEFAULT_MODEL_PATH, help="Path to sban_weighted_stacking_model.joblib", ) parser.add_argument( "--input", type=Path, required=True, help="Input parquet with representation columns", ) parser.add_argument( "--output", type=Path, required=True, help="Output parquet with prediction and prob_* columns", ) parser.add_argument( "--evaluate", action="store_true", help="Print metrics when dataset_name column is present", ) return parser.parse_args() def main() -> None: args = parse_args() if not args.model_path.is_file(): raise FileNotFoundError(f"Model not found: {args.model_path}") if not args.input.is_file(): raise FileNotFoundError(f"Input not found: {args.input}") predictor = load_predictor(args.model_path) df = pd.read_parquet(args.input) predictions = predictor.predict(df) args.output.parent.mkdir(parents=True, exist_ok=True) predictions.to_parquet(args.output, index=False) print(f"Wrote {len(predictions)} rows to {args.output}") if args.evaluate: evaluate_predictions(predictions) if __name__ == "__main__": main()