prajwalpujari16

EdgeSLM-qwen0.5b-v2

README

License: apache-2.0

Labels

Table with columns: Label, Definition used in the prompt
Label	Definition used in the prompt
`PUBLIC`	approved information intended for external sharing
`INTERNAL`	routine company-only procedures or workflows
`CONFIDENTIAL`	private business, employee, customer, or contract information
`RESTRICTED`	credentials, keys, tokens, passwords, or access-bypass instructions

How to use

The model is not used for generation. It reads the logits of the first token of each label after a fixed prompt, in a single forward pass:

python

import torchfrom peft import PeftModelfrom transformers import AutoModelForCausalLM, AutoTokenizer BASE = "Qwen/Qwen2.5-0.5B-Instruct"ADAPTER = "prajwalpujari16/EdgeSLM-qwen0.5b-v2"LABELS = ("PUBLIC", "INTERNAL", "CONFIDENTIAL", "RESTRICTED")GUIDANCE = (    "PUBLIC = approved information intended for external sharing.\n"    "INTERNAL = routine company-only procedures or workflows.\n"    "CONFIDENTIAL = private business, employee, customer, or contract information.\n"    "RESTRICTED = credentials, keys, tokens, passwords, or access-bypass instructions.") tokenizer = AutoTokenizer.from_pretrained(BASE)model = PeftModel.from_pretrained(AutoModelForCausalLM.from_pretrained(BASE), ADAPTER).eval()label_ids = [tokenizer(" " + label, add_special_tokens=False)["input_ids"][0] for label in LABELS] def classify(text):    prompt = (        "Classify the sensitivity of the following company document.\n\n"        f"Use these definitions:\n{GUIDANCE}\n\n"        f"Document:\n{text}\n\n"        "Classification:"    )    inputs = tokenizer(prompt, return_tensors="pt")    with torch.no_grad():        logits = model(**inputs).logits[0, -1]    probs = torch.softmax(logits[label_ids].float(), dim=0)    return dict(zip(LABELS, probs.tolist())) print(classify("Use svc_ingest / Copper!2291 for the ledger service until Monday."))

The prompt has to fit in 512 tokens with Classification: at the end. For longer documents, split the text into chunks; the repository's src/inference/chunking.py does this, and src/scan.py classifies whole folders of PDF, Word, email and text files.

Results

Accuracy with 95% Wilson intervals. Everything below except the last table is on short (one or two sentence) documents written for the project.

Table with columns: Set, n, This adapter, TF-IDF baseline, same training data
Set	n	This adapter	TF-IDF baseline, same training data
`contrast`: 40 pairs sharing vocabulary but differing in label	80	88.8% [80–94], 31/40 pairs fully right	75.0% [65–83], 21/40
`reviewed-heldout`: hand-written, independently re-labelled	60	90.0% [80–95]	95.0% [86–98]
`realistic`	40	90.0%	97.5%

The gap over TF-IDF on contrast is significant (exact McNemar, p = 0.043). The gap on reviewed-heldout is not (p = 0.38). The contrast set and the training data generator were written by the same author, so contrast scores are likely optimistic.

On 40 real documents (20 Enron emails, 20 GitLab Handbook pages):

Table with columns: System, Agreement
System	Agreement
Always answer `INTERNAL`	26/40
TF-IDF, same training data	19/40
This adapter	11/40

19 of the 26 INTERNAL documents and 7 of the 9 PUBLIC ones were called CONFIDENTIAL or RESTRICTED, mostly with over 97% confidence. Real email is full of names, phone numbers and addresses, which the synthetic training data mostly placed in CONFIDENTIAL examples. These reference labels were assigned by Claude, not a person, and only 5 of the 40 documents are CONFIDENTIAL or RESTRICTED.

Edge deployment

Merged into the base model and exported to ONNX with 4-bit weights (MatMulNBits, block size 32, activations kept in fp32), it is a 392 MB file that classifies a document in a median 469 ms using about 1 GB of RAM on a Ryzen 5 6600H laptop CPU. Its accuracy matches full precision on every test set above. Dynamic INT8, which also quantises activations, did not hold up in this project and should be avoided.

Training

Data: 2,360 synthetic documents (data/splits/train.jsonl, data/raw/boundary_training.jsonl and data/v2/train.jsonl in the repository). The v2 generator writes live secrets with no keyword, public docs full of security vocabulary, and placeholder-versus-real-value pairs.
Objective: loss on the label tokens only (prompt/completion SFT with TRL).
LoRA: rank 16, alpha 32, dropout 0.05, on all linear layers (q, k, v, o, gate, up, down projections).
Schedule: 3 epochs, learning rate 2e-4 with cosine decay, effective batch size 16, fp16. About 21 minutes on an RTX 3050 (6 GB). The best checkpoint by validation loss is the one published.

Intended use and limitations

Intended for research and teaching: studying how small language models classify document sensitivity, and how they fail. It is not suitable for deciding what may leave an organisation.

Trained only on synthetic sentences; fails to transfer to real documents (above).
Weak on adversarial phrasing; catches only 4 of 10 credential-bearing documents in adversarial.
English only.
No human has verified any evaluation label.

License

The adapter is released under Apache 2.0, the license of the base model. The repository's code is MIT.

Available on FriendliAI

Dedicated Endpoints

Run this model inference on single tenant GPU with unmatched speed and reliability at scale.

Learn more

Container

Run this model inference with full control and performance in your environment.

Learn more

Model Details

Model Provider

prajwalpujari16

Model Tree

Base

Qwen/Qwen2.5-0.5B-Instruct

Adapter

this model

Input Modalities

Text

Output Modalities