canalan/MalwareDatasetClassification

MalwareDatasetClassification (SBAN)

[Türkçe dokümantasyon](README.tr.md)

Multiclass pipeline for malware dataset origin classification on SBAN: four synchronized text views per sample → predict which sub-corpus it belongs to (bodmas, dike, malwarebazaar, sorel20m).

This repository contains code, notebooks, and `sban_weighted_stacking_model.joblib`. No SBAN parquet or raw JSON is distributed; obtain SBAN separately.

Task and labels

Data preparation pipeline (scripts 01–11)

End-to-end flow on local SBAN exports:

1. 1.`01_make_a_dataframe.py` — Merge JSON shards under data/M1/SBAN-MA-JUN25 into SBAN.parquet (four representations aligned by ID).
2. 2.`02_validate_data.py` — Schema, missing values, duplicates, cross-dataset ID overlap, content fingerprints (see Data quality).
3. 3.`03_make_clean_dataframe.py` — Cleaning rules → SBAN_clean.parquet.
4. 4.`04_analyze_prompt_residue.py` — Count LLM/prompt boilerplate phrases per representation (Prompt residue).
5. 5.`05_split_dataframe.py` — Stratified train / validation / test parquet files.
6. 6.`06_make_features.py` — Optional TF-IDF .npz features for alternate experiments.
7. 7.`07`–`11` — Per-representation audits and source cleaning (08_clean_source_code.py uses 07_audit_source_code.py).

Notebooks:

• —`baseline.ipynb` — Early fusion / baseline stacking comparisons.
• —`svc_sban.ipynb` — Production model: per-representation TF-IDF + numeric features, ID-based feature pruning, class-weight search, weighted LinearSVC bases, HistGradientBoostingClassifier meta learner, joblib export.
• —`inference.ipynb` — Load exported bundle; validation/test metrics; synthetic demo row.

Canonical runtime entrypoint: `inference.py` (CLI + StackingPredictor).

Data quality findings

Summaries below come from running the numbered scripts on the full merged SBAN table (before train/val/test split). Reproduce with your own copy of the data.

Cross-dataset ID overlap (content match rate)

Shared IDs across corpus pairs; percentages = share of common IDs where that column’s text is byte-identical (02_validate_data.py, section 9).

High overlap for bodmas × sorel20m (especially source) motivates careful splitting and explains why the classifier must use subtle cues, not only exact string identity across corpora.

Rows with four aligned representations

After merge / alignment (01_make_a_dataframe.py):

“Matched” = samples where all four representation fields are present for labeling and training.

Prompt residue analysis

04_analyze_prompt_residue.py scans fixed English phrases (e.g. “your code”, “here”, “additional”) across columns. Illustrative totals on cleaned data:

Most residue sits in source and NLD; source cleaning scripts (07/08) target audit failures before modeling.

Model architecture

Artifact: `sban_weighted_stacking_model.joblib` (bundle_version: 1, trained with scikit-learn 1.6.1).

For each r ∈ {asm, binary, source, nld}:
  text → TF-IDF (binary: hex → byte tokens + instsep)
       + 6 numeric stats (length, tokens, entropy, …)
       → StandardScaler
       → sparse hstack → column subset (selected_indices from ID pruning)
       → LinearSVC (tuned class weights) → decision_function (4 scores)

Meta:
  hstack(all base decision scores + all scaled numeric blocks)
  → HistGradientBoostingClassifier
  → class probabilities

Bundle keys: representation_order, representation_columns, numeric_feature_names, label_encoder, meta_model, representations (vectorizer, scaler, indices, base model), selected_class_weight_configs, metadata.

Training details and ablations: `svc_sban.ipynb`.

Feature pruning (TF-IDF columns)

Implemented in `svc_sban.ipynb` (cells after the first per-representation LinearSVC bases):

1. 1.Importance — For each representation, mean |coef_| over classes from final_base_models (TF-IDF tokens + six numeric stats).
2. 2.Sort ascending — Lowest-importance names are dropped first.
3. 3.Ratio sweep — Validation macro-F1 was plotted for many removal ratios (roughly 5–60% and 65–80% in the analysis figures); the exported model uses a single setting.
4. 4.Production choice — `feature_pruning_ratio = 0.65`: remove the lowest 65% of the ranked feature list for TF-IDF vocabulary entries. The six numeric columns (char_count, line_count, token_count, avg_line_length, unique_token_ratio, char_entropy) are always kept and re-appended via fixed column indices after TF-IDF subsetting.

Validation macro-F1 at 65% feature removal (same notebook run):

These pruned column sets are stored in the joblib bundle as representations[r]["selected_indices"] (feature step only; ID pruning below may reuse the same index vector).

ID pruning (training samples)

Overlapping bodmas vs sorel20m IDs motivate dropping ambiguous training rows before refitting bases:

1. 1.Fix feature pruning at 65% and fit a temporary LinearSVC on pruned features.
2. 2.Score each training row in `bodmas` and `sorel20m` only: sparse TF-IDF presence (binary) dotted with pruned-model TF-IDF coefficient magnitudes → importance_score.
3. 3.Grid — For each representation, remove the lowest-scoring `id_prune_ratios` fraction per class (5%, 10%, …, 70%), refit on remaining train rows, measure validation macro-F1 → id_pruning_summary in the notebook.
4. 4.Production choice — selected_id_prune_ratios:

Final stacking retrains ID-pruned bases (5-fold OOF decision scores), then class-weight search and meta learner on top of that pipeline. `dike` and `malwarebazaar` rows are never removed by this step.

Split sizes used in training notebook

(Test counts from inference.ipynb evaluation on exported bundle.)

Base models on validation (svc_sban.ipynb)

Single-representation LinearSVC decision scores, validation set:

Source is the strongest single view; nld alone is weakest but adds complementary signal in the stack.

Final exported model — validation & test

Metrics from `inference.ipynb` with sban_weighted_stacking_model.joblib (matches weighted meta validation in svc_sban.ipynb before export).

Validation (n = 16,167)

Test (n = 32,334)

Minority classes (dike, malwarebazaar) remain the hardest; weighted class tuning in svc_sban.ipynb targets that imbalance.

Inference schema

Installation

pip install -r requirements-inference.txt   # predict only
pip install -r requirements.txt             # full pipeline + notebooks

Use scikit-learn 1.6.1 when loading the joblib bundle.

Running inference

python inference.py \
  --model-path sban_weighted_stacking_model.joblib \
  --input /path/to/SBAN_test.parquet \
  --output predictions.parquet \
  --evaluate

from inference import load_predictor
import pandas as pd

predictor = load_predictor("sban_weighted_stacking_model.joblib")
out = predictor.predict(pd.read_parquet("/path/to/samples.parquet"))

Notebook: inference.ipynb — Colab or local setup → demo row → validation/test cells (update parquet paths).

Reproducing the production model

1. 1.Obtain SBAN and build parquets via 01–05 (and cleaning/audit scripts as needed).
2. 2.Open `svc_sban.ipynb` (Colab or local), point to SBAN_train/val/test.parquet.
3. 3.Run training cells; export `sban_weighted_stacking_model.joblib` to the repo root.
4. 4.Verify with `inference.py` or `inference.ipynb`.

Citation and security

• —Cite the SBAN dataset authors; this repo does not redistribute their files.
• —`joblib.load` uses pickle — only load bundles from this project or your own exports.