Views
No views yet
NONE if
there was no relationship between the entities. The model implements a classifier head
on top of ModernBERT-large in order to take advantage of the extended context window.NONE (No relationship between entities)CAUSESWORSENSIMPROVESRELATES_TOEXPERIENCESTRIGGERShf download dzur658/TherapyBERT-RE-001 --local-dir .1import torch
2from transformers import AutoTokenizer
3
4# make sure the classification head is present before importing!
5from modern_bert_re_layers import ModernBERT_Entity_Pooling_RE
6
7from typing import Dict, List, Optional, Sequence, Tuple, Union
8import re
9
10UNIQUE_LABELS = ["NONE", "CAUSES", "WORSENS", "IMPROVES", "RELATES_TO", "EXPERIENCES", "TRIGGERS"]
11
12# here we load the base tokenizer, but since we added 4 new tokens
13# for entity extraction we will have to manually add them below
14TOKENIZER_MODEL = "answerdotai/ModernBERT-large"
15tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL)
16
17# add special tokens for entity extraction
18SPECIAL_TOKENS = {"additional_special_tokens": ["[E1]", "[/E1]", "[E2]", "[/E2]"]}
19tokenizer.add_special_tokens(SPECIAL_TOKENS)
20
21# use on cuda, mps, or cpu
22# for example cuda
23device = torch.device("cuda")
24
25# set path to weights/config
26MODEL_PATH = "./therapy-modernbert-re-final"
27
28model = ModernBERT_Entity_Pooling_RE.from_checkpoint(
29 MODEL_PATH,
30 tokenizer=tokenizer,
31 map_location=device,
32 )
33
34# set to eval mode for inference
35model.eval()
36
37# cast bf16 to cuda/mps, for cpu use fp32 operations
38model.to(device, dtype=torch.bfloat16 if device.type in ["cuda", "mps"] else torch.float32)
39
40# example input
41test_cases = [
42 (
43 "My anxiety has been getting worse since the calls from my ex-husband started again.",
44 "ex-husband",
45 "anxiety",
46 ),
47 (
48 "I avoid crowded stores because the flashing lights can trigger a panic attack.",
49 "flashing lights",
50 "panic attack",
51 ),
52 ]
53
54# helper function to extract entities from the text
55def mark_entities(text: str, source: str, target: str) -> str:
56 if not source or not target:
57 raise ValueError("Both source and target entity strings are required.")
58
59 marked_text = re.sub(f"({re.escape(source)})", r"[E1]\1[/E1]", text, count=1)
60 marked_text = re.sub(f"({re.escape(target)})", r"[E2]\1[/E2]", marked_text, count=1)
61
62 if "[E1]" not in marked_text or "[E2]" not in marked_text:
63 raise ValueError("Could not find both entities in the provided text.")
64
65 return marked_text
66
67# inference function
68def predict_marked_text(marked_text: str, top_k: Optional[int] = 1) -> Union[Dict[str, float], List[Dict[str, float]]]:
69 # supports full 8192 token context window
70 max_len = 8192
71
72 inputs = tokenizer(
73 marked_text,
74 return_tensors="pt",
75 truncation=True,
76 max_length=max_len,
77 )
78 inputs = {key: value.to(device) for key, value in inputs.items()}
79
80 with torch.no_grad():
81 outputs = model(**inputs)
82 logits = outputs["logits"]
83 probabilities = torch.softmax(logits, dim=-1)[0]
84
85 scores, indices = torch.sort(probabilities, descending=True)
86 predictions = [
87 {
88 "label": model.id2label[int(index)],
89 "score": float(score),
90 }
91 for score, index in zip(scores.tolist(), indices.tolist())
92 ]
93
94 if top_k is None:
95 return predictions
96
97 limited_predictions = predictions[:top_k ]
98
99 if top_k == 1:
100 return limited_predictions
101 else:
102 return limited_predictions
103
104for case in test_cases:
105 print(case)
106 text, source, target = case
107 marked_text = mark_entities(text, source, target)
108
109 # top k can be passed here otherwise defaults to 1
110 predictions = predict_marked_text(marked_text, top_k=3)
111 for pred in predictions:
112 print(f" -> {pred['label']}: {pred['score']:.4f}")
113 print("-" * 80)('My anxiety has been getting worse since the calls from my ex-husband started again.', 'ex-husband', 'anxiety')
-> IMPROVES: 0.2754
-> NONE: 0.2246
-> TRIGGERS: 0.1777
--------------------------------------------------------------------------------
('I avoid crowded stores because the flashing lights can trigger a panic attack.', 'flashing lights', 'panic attack')
-> TRIGGERS: 0.4883
-> NONE: 0.1895
-> CAUSES: 0.1162
--------------------------------------------------------------------------------1e-3)eval_loss: 0.9485eval_accuracy: 0.6974eval_f1_macro: 0.2671@misc{TherapyBERT RE,
title = {TherapyBERT RE 001},
author = {{Alex Dzurec}},
month = {March},
year = {2026},
url = {https://huggingface.co/dzur658/TherapyBERT-RE-001}
}