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
MalwareDatasetClassification (SBAN)
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.
| Resource | Location |
|---|---|
| Source | github.com/berkecanalan/MalwareDatasetClassification |
| Weights | Repo root + huggingface.co/canalan/MalwareDatasetClassification |
| License | Apache-2.0 (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:
01_make_a_dataframe.py— Merge JSON shards underdata/M1/SBAN-MA-JUN25intoSBAN.parquet(four representations aligned byID).02_validate_data.py— Schema, missing values, duplicates, cross-datasetIDoverlap, content fingerprints (see Data quality).03_make_clean_dataframe.py— Cleaning rules →SBAN_clean.parquet.04_analyze_prompt_residue.py— Count LLM/prompt boilerplate phrases per representation (Prompt residue).05_split_dataframe.py— Stratified train / validation / test parquet files.06_make_features.py— Optional TF-IDF.npzfeatures for alternate experiments.07–11— Per-representation audits and source cleaning (08_clean_source_code.pyuses07_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, weightedLinearSVCbases,HistGradientBoostingClassifiermeta 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).
| 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).
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):
- Importance — For each representation, mean
|coef_|over classes fromfinal_base_models(TF-IDF tokens + six numeric stats). - Sort ascending — Lowest-importance names are dropped first.
- 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.
- 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:
- Fix feature pruning at 65% and fit a temporary
LinearSVCon pruned features. - Score each training row in
bodmasandsorel20monly: sparse TF-IDF presence (binary) dotted with pruned-model TF-IDF coefficient magnitudes →importance_score. - Grid — For each representation, remove the lowest-scoring
id_prune_ratiosfraction per class (5%, 10%, …, 70%), refit on remaining train rows, measure validation macro-F1 →id_pruning_summaryin the notebook. - 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
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
- Obtain SBAN and build parquets via
01–05(and cleaning/audit scripts as needed). - Open
svc_sban.ipynb(Colab or local), point toSBAN_train/val/test.parquet. - Run training cells; export
sban_weighted_stacking_model.joblibto the repo root. - Verify with
inference.pyorinference.ipynb.
Citation and security
- Cite the SBAN dataset authors; this repo does not redistribute their files.
joblib.loaduses pickle — only load bundles from this project or your own exports.
- Downloads last month
- -