Instructions to use canalan/MalwareDatasetClassification with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- Scikit-learn
How to use canalan/MalwareDatasetClassification with Scikit-learn:
from huggingface_hub import hf_hub_download import joblib model = joblib.load( hf_hub_download("canalan/MalwareDatasetClassification", "sklearn_model.joblib") ) # only load pickle files from sources you trust # read more about it here https://skops.readthedocs.io/en/stable/persistence.html - Notebooks
- Google Colab
- Kaggle
| license: apache-2.0 | |
| library_name: sklearn | |
| tags: | |
| - malware | |
| - multi-class-classification | |
| - stacking | |
| - sban | |
| pipeline_tag: text-classification | |
| # MalwareDatasetClassification (SBAN) | |
| **[TΓΌrkΓ§e dokΓΌmantasyon](https://github.com/berkecanalan/MalwareDatasetClassification/blob/main/README.tr.md)** (Hubβda yalnΔ±zca Δ°ngilizce `README.md`; TΓΌrkΓ§e metin GitHubβda) | |
| Multiclass pipeline for **malware dataset origin classification** on [SBAN](https://github.com/ma-soreto/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. | |
| | Resource | Location | | |
| |----------|----------| | |
| | Source | [github.com/berkecanalan/MalwareDatasetClassification](https://github.com/berkecanalan/MalwareDatasetClassification) | | |
| | Weights | Repo root + [huggingface.co/canalan/MalwareDatasetClassification](https://huggingface.co/canalan/MalwareDatasetClassification) | | |
| | License | Apache-2.0 ([LICENSE](LICENSE)) | | |
| --- | |
| ## Task and labels | |
| | | | | |
| |---|---| | |
| | **Input** | `assembly_code`, `binary_code`, `source_code`, `NLD` for one sample | | |
| | **Output** | `dataset_name` β {`bodmas`, `dike`, `malwarebazaar`, `sorel20m`} | | |
| | **Scope** | Dataset **provenance** classification, not generic malware detection | | |
| --- | |
| ## Data preparation pipeline (scripts `01`β`11`) | |
| End-to-end flow on local SBAN exports: | |
| 1. **`01_make_a_dataframe.py`** β Merge JSON shards under `data/M1/SBAN-MA-JUN25` into `SBAN.parquet` (four representations aligned by `ID`). | |
| 2. **`02_validate_data.py`** β Schema, missing values, duplicates, cross-dataset `ID` overlap, content fingerprints (see [Data quality](#data-quality-findings)). | |
| 3. **`03_make_clean_dataframe.py`** β Cleaning rules β `SBAN_clean.parquet`. | |
| 4. **`04_analyze_prompt_residue.py`** β Count LLM/prompt boilerplate phrases per representation ([Prompt residue](#prompt-residue-analysis)). | |
| 5. **`05_split_dataframe.py`** β Stratified train / validation / test parquet files. | |
| 6. **`06_make_features.py`** β Optional TF-IDF `.npz` features for alternate experiments. | |
| 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 `ID`s across corpus pairs; percentages = share of common IDs where that columnβs text is **byte-identical** (`02_validate_data.py`, section 9). | |
| | Pair | Common IDs | assembly | binary | source | | |
| |------|------------|----------|--------|--------| | |
| | bodmas Γ sorel20m | 806 | 72% | 74% | 91% | | |
| | bodmas Γ malwarebazaar | 520 | 27% | 25% | 0% | | |
| | bodmas Γ dike | 201 | 30% | 24% | 0% | | |
| | dike Γ malwarebazaar | 82 | 23% | 22% | 4% | | |
| | malwarebazaar Γ sorel20m | 99 | 41% | 47% | 0% | | |
| | dike Γ sorel20m | 46 | 44% | 48% | 0% | | |
| 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`): | |
| | Dataset | Rows | Matched (4 repr.) | | |
| |---------|------|-------------------| | |
| | bodmas | 82,032 | 757 | | |
| | dike | 5,342 | 669 | | |
| | malwarebazaar | 6,048 | 905 | | |
| | sorel20m | 71,319 | 726 | | |
| β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: | |
| | Phrase | assembly | binary | source | NLD | Total | | |
| |--------|----------|--------|--------|-----|-------| | |
| | your code | 1 | 0 | 572 | 1 | 573 | | |
| | add main function | 0 | 0 | 69 | 0 | 69 | | |
| | code goes | 0 | 0 | 140 | 0 | 140 | | |
| | implementation goes | 0 | 0 | 49 | 0 | 49 | | |
| | corrected | 0 | 0 | 168 | 16 | 169 | | |
| | here | 128 | 0 | 1,493 | 323 | 1,798 | | |
| | no comments | 0 | 0 | 53 | 0 | 53 | | |
| | additional | 37 | 0 | 76 | 1,103 | 1,185 | | |
| 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**). | |
| ```text | |
| 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. **Importance** β For each representation, mean `|coef_|` over classes from `final_base_models` (TF-IDF tokens + six numeric stats). | |
| 2. **Sort ascending** β Lowest-importance names are dropped first. | |
| 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. **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): | |
| | Representation | Macro F1 (val) | Columns after prune | | |
| |----------------|----------------|---------------------| | |
| | asm | 0.7009 | 26,259 | | |
| | binary | 0.6609 | 26,259 | | |
| | nld | 0.4870 | 26,259 | | |
| | source | 0.8802 | 26,257 | | |
| 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. Fix **feature pruning at 65%** and fit a temporary `LinearSVC` on pruned features. | |
| 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. **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. **Production choice** β `selected_id_prune_ratios`: | |
| | Representation | ID remove ratio | Val macro F1 | Train rows kept | Removed bodmas / sorel20m | | |
| |----------------|-----------------|--------------|-----------------|---------------------------| | |
| | asm | **10%** | 0.6978 | 102,538 | 5,689 / 4,941 | | |
| | binary | **5%** | 0.6625 | 107,854 | 2,844 / 2,470 | | |
| | source | **40%** | 0.8712 | 70,646 | 22,756 / 19,766 | | |
| | nld | **40%** | 0.4804 | 70,646 | 22,756 / 19,766 | | |
| 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 | |
| | Split | Rows | bodmas | dike | malwarebazaar | sorel20m | | |
| |-------|------|--------|------|---------------|----------| | |
| | Train | 113,168 | 56,892 | 3,267 | 3,594 | 49,415 | | |
| | Validation | 16,167 | 8,128 | 467 | 513 | 7,059 | | |
| | Test | 32,334 | 16,255 | 933 | 1,027 | 14,119 | | |
| (Test counts from `inference.ipynb` evaluation on exported bundle.) | |
| ### Base models on validation (`svc_sban.ipynb`) | |
| Single-representation `LinearSVC` decision scores, validation set: | |
| | Representation | Accuracy | Macro F1 | Weighted F1 | | |
| |----------------|----------|----------|---------------| | |
| | asm | 0.8983 | 0.7099 | 0.8902 | | |
| | binary | 0.8426 | 0.6607 | 0.8345 | | |
| | source | 0.9253 | 0.8811 | 0.9251 | | |
| | nld | 0.6621 | 0.4997 | 0.6539 | | |
| **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)** | |
| | | Accuracy | Macro F1 | Weighted F1 | | |
| |---|----------|----------|-------------| | |
| | Overall | **0.9413** | **0.9097** | **0.9412** | | |
| | Class | Precision | Recall | F1 | Support | | |
| |-------|-----------|--------|-----|---------| | |
| | bodmas | 0.9551 | 0.9398 | 0.9474 | 8,128 | | |
| | dike | 0.9125 | 0.8266 | 0.8674 | 467 | | |
| | malwarebazaar | 0.8986 | 0.8635 | 0.8807 | 513 | | |
| | sorel20m | 0.9306 | 0.9562 | 0.9433 | 7,059 | | |
| **Test (n = 32,334)** | |
| | | Accuracy | Macro F1 | Weighted F1 | | |
| |---|----------|----------|-------------| | |
| | Overall | **0.9379** | **0.9012** | **0.9378** | | |
| | Class | Precision | Recall | F1 | Support | | |
| |-------|-----------|--------|-----|---------| | |
| | bodmas | 0.9532 | 0.9364 | 0.9447 | 16,255 | | |
| | dike | 0.8909 | 0.8489 | 0.8694 | 933 | | |
| | malwarebazaar | 0.8885 | 0.8150 | 0.8502 | 1,027 | | |
| | sorel20m | 0.9272 | 0.9545 | 0.9407 | 14,119 | | |
| Minority classes (`dike`, `malwarebazaar`) remain the hardest; weighted class tuning in `svc_sban.ipynb` targets that imbalance. | |
| --- | |
| ## Inference schema | |
| | Column | Required for predict | Notes | | |
| |--------|-------------------|--------| | |
| | `assembly_code`, `binary_code`, `source_code`, `NLD` | Yes | | | |
| | `ID` | No | Preserved in output | | |
| | `dataset_name` | No | For `--evaluate` / notebook metrics | | |
| --- | |
| ## Installation | |
| ```bash | |
| 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 | |
| ```bash | |
| python inference.py \ | |
| --model-path sban_weighted_stacking_model.joblib \ | |
| --input /path/to/SBAN_test.parquet \ | |
| --output predictions.parquet \ | |
| --evaluate | |
| ``` | |
| ```python | |
| 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. Obtain SBAN and build parquets via `01`β`05` (and cleaning/audit scripts as needed). | |
| 2. Open **`svc_sban.ipynb`** (Colab or local), point to `SBAN_train/val/test.parquet`. | |
| 3. Run training cells; export **`sban_weighted_stacking_model.joblib`** to the repo root. | |
| 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. | |