mentalkg-xlmr-node

Multi-label concept extractor for mental-health journal entries. Fine-tuned XLM-RoBERTa base with a 130-way sigmoid head. Predicts which concepts from a fixed 130-label vocabulary are expressed in a short journal text (English or German).

Paired with mentalkg-xlmr-edge to build a full journal-to-graph pipeline (nodes from this model, connectivity from the edge model, relation types from a lookup table).

Usage

import json, torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

tok = AutoTokenizer.from_pretrained("Niklas1102/mentalkg-xlmr-node")
model = AutoModelForSequenceClassification.from_pretrained("Niklas1102/mentalkg-xlmr-node").eval()

meta = json.loads(open(tok.name_or_path + "/meta.json").read())  # or hf_hub_download
labels = json.loads(open(tok.name_or_path + "/labels.json").read())["index_to_label"]

text = "I barely slept and the deadline is tomorrow."
enc = tok(text, return_tensors="pt", truncation=True, max_length=256)
with torch.no_grad():
    probs = torch.sigmoid(model(**enc).logits[0])
predicted = [labels[i] for i, p in enumerate(probs) if p >= meta["threshold"]]
print(predicted)

Input

One journal entry as a raw string, up to 256 tokens (XLM-R tokenizer). Longer text is right-truncated during training and at inference.

Output

A 130-length logit vector. Apply sigmoid, then keep labels above meta["threshold"] (0.28). The 130 concept ids and their surface labels are in labels.json.

Results on the 8,246-example test split

slice	micro-F1	macro-F1	micro-P	micro-R
combined EN+DE	0.911	0.849	0.920	0.902
EN test slice (combined model)	0.927	n/a	n/a	n/a
DE test slice (combined model)	0.894	n/a	n/a	n/a

Reproducibility: participant-independent 50/30/20 split stratified by language, seed 42. Config: --head-bias-init, lr 3e-5, 9 epochs, plain BCE loss.

Training data

Synthetic bilingual (EN/DE) journal + graph corpus, 41,315 accepted samples across 3,410 participants. Journal texts generated by openai/gpt-4o-mini from ground-truth graphs; each entry LLM-verified for node coverage, edge expression, temporal correctness, and hallucination. Data card: mentalkg. Full training protocol in the code repo.

Intended use

Research on narrative graph extraction, mental-health text mining benchmarks, and downstream tooling that reasons over concept co-occurrence in short journal entries.

Out of scope

Diagnostic use, clinical decision support, real-time triage. The training data is LLM-generated fiction, and outputs describe narrated content rather than the writer's clinical state. The label vocabulary is a fixed 130-item research inventory, separate from any validated clinical ontology.

Ethics

Training data is synthetic. No human subjects, no scraped text, no personally identifying content. The German entries were generated in a native register and idiomatically verified; the English entries were generated in the same pipeline. Both languages use identical guardrails against connector words, clinical labels, self-harm and violence.

License

MIT.