Views
No views yet
BODY_PART: Anatomical regions, organs, or physiological systems (e.g. 十二指腸, 乳房, 角膜, 腎臟).SURGERY_TYPE: Surgical actions, techniques, or procedure types (e.g. 切除術, 縫合術, 移植, 置換術).TECH_DEVICE: Medical devices, surgical technologies, or specific equipment (e.g. 達文西, 雷射, 腹腔鏡, 超音波乳化).model.safetensors: Fine-tuned weights.config.json: Architecture, hyperparameters, and label definitions (id2label / label2id).tokenizer.json: Serialized tokenizer configuration.tokenizer_config.json: Instantiate arguments.1import os
2import torch
3from transformers import AutoTokenizer, AutoModelForTokenClassification
4
5# Model folder (local path containing these files)
6MODEL_DIR = os.path.dirname(os.path.abspath(__file__))
7
8LABEL_LIST = ["O", "B-BODY_PART", "I-BODY_PART", "B-SURGERY_TYPE", "I-SURGERY_TYPE", "B-TECH_DEVICE", "I-TECH_DEVICE"]
9
10def extract_entities(text: str, tokenizer, model) -> list:
11 """Tokenizes text and groups token classifications into character-aligned entity spans."""
12 if not text.strip():
13 return []
14
15 inputs = tokenizer(
16 text,
17 return_offsets_mapping=True,
18 return_tensors="pt",
19 truncation=True,
20 max_length=512
21 )
22
23 device = next(model.parameters()).device
24 inputs = {k: v.to(device) for k, v in inputs.items()}
25
26 with torch.no_grad():
27 outputs = model(**{k: v for k, v in inputs.items() if k != "offset_mapping"})
28
29 logits = outputs.logits
30 predictions = torch.argmax(logits, dim=2)[0].cpu().numpy()
31 offsets = inputs["offset_mapping"][0].cpu().numpy()
32
33 entities = []
34 current_entity = None
35
36 for idx, offset in enumerate(offsets):
37 start, end = offset
38 # Ignore padding and special tokens
39 if start == 0 and end == 0:
40 continue
41
42 label = LABEL_LIST[predictions[idx]]
43
44 if label.startswith("B-"):
45 if current_entity:
46 entities.append(current_entity)
47 entity_type = label.split("-")[1]
48 current_entity = {
49 "label": entity_type,
50 "start": int(start),
51 "end": int(end),
52 "text": text[start:end]
53 }
54 elif label.startswith("I-"):
55 entity_type = label.split("-")[1]
56 if current_entity and current_entity["label"] == entity_type:
57 current_entity["end"] = int(end)
58 current_entity["text"] = text[current_entity["start"]:int(end)]
59 else:
60 if current_entity:
61 entities.append(current_entity)
62 current_entity = {
63 "label": entity_type,
64 "start": int(start),
65 "end": int(end),
66 "text": text[start:end]
67 }
68 else: # 'O'
69 if current_entity:
70 entities.append(current_entity)
71 current_entity = None
72
73 if current_entity:
74 entities.append(current_entity)
75
76 return entities
77
78def main():
79 print(f"Loading model from {MODEL_DIR}...")
80 tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
81 model = AutoModelForTokenClassification.from_pretrained(MODEL_DIR)
82
83 # Detect Apple Silicon GPU (mps), CUDA, or CPU
84 device = "cpu"
85 if torch.backends.mps.is_available():
86 device = "mps"
87 elif torch.cuda.is_available():
88 device = "cuda"
89 model = model.to(device)
90 model.eval()
91
92 # Example Inference
93 test_sentence = "病患因右側腹股溝疝氣住院，在門診實施了微創內視鏡疝氣修補術。"
94 entities = extract_entities(test_sentence, tokenizer, model)
95
96 print(f"\nInput: {test_sentence}")
97 print("Extracted Entities:")
98 for ent in entities:
99 print(f" - {ent['text']} | Label: {ent['label']} | Spans: ({ent['start']}, {ent['end']})")
100
101if __name__ == "__main__":
102 main()pip install torch transformers