478_生物医学NLP_BioNER¶

# ---- 1. 使用 BioBERT 进行生物医学NER ----
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

# 加载专门用于基因/疾病/化合物NER的BioBERT微调模型
ner_model_name = "pruas/BENT-PubMedBERT-NER-Gene"  # 基因NER模型

tokenizer = AutoTokenizer.from_pretrained(ner_model_name)
model = AutoModelForTokenClassification.from_pretrained(ner_model_name)

# 创建NER流水线
bio_ner = pipeline(
    "ner",
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="simple"  # 自动合并跨subtoken的实体
)

# 生物医学文本示例
bio_text = """
The BRCA1 gene is associated with increased risk of breast cancer.
Tamoxifen is commonly used to treat estrogen receptor-positive breast cancer.
TP53 mutations are found in approximately 50% of human cancers.
"""

entities = bio_ner(bio_text)
print("识别到的生物医学实体:")
for ent in entities:
    print(f"  [{ent['entity_group']}] {ent['word']} (置信度: {ent['score']:.3f})")

# ---- 2. 使用 BERN2 API 进行多类型BioNER ----
import requests

def bern2_ner(text):
    """调用 BERN2 在线API（多类实体：基因/疾病/化合物/物种/突变）"""
    url = "http://bern2.korea.ac.kr/plain"
    response = requests.post(url, json={"text": text})
    result = response.json()
    for ann in result.get('annotations', []):
        entity_text = text[ann['span']['begin']:ann['span']['end']]
        entity_types = [t['name'] for t in ann.get('obj', [])]
        print(f"  {entity_text}: {entity_types}")

# ---- 3. PubMedBERT 微调 BioNER（自定义数据集）----
from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import TrainingArguments, Trainer
from datasets import load_dataset
import numpy as np
from seqeval.metrics import f1_score

# 加载BC5CDR数据集（化合物+疾病NER的标准benchmark）
dataset = load_dataset("ncats/bc5cdr")

# BC5CDR标签
label_list = dataset["train"].features["ner_tags"].feature.names
label2id = {l: i for i, l in enumerate(label_list)}
id2label = {i: l for i, l in enumerate(label_list)}

# 加载PubMedBERT（专域预训练，无通用混合）
model_checkpoint = "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
model = AutoModelForTokenClassification.from_pretrained(
    model_checkpoint,
    num_labels=len(label_list),
    id2label=id2label,
    label2id=label2id
)

def tokenize_and_align(examples):
    """分词并对齐NER标签（处理wordpiece切割）"""
    tokenized = tokenizer(
        examples["tokens"],
        is_split_into_words=True,   # 输入已分词
        truncation=True,
        max_length=512
    )
    labels = []
    for i, label_ids in enumerate(examples["ner_tags"]):
        word_ids = tokenized.word_ids(batch_index=i)
        previous_word_idx = None
        label_list_aligned = []
        for word_idx in word_ids:
            if word_idx is None:
                label_list_aligned.append(-100)  # 特殊token忽略
            elif word_idx != previous_word_idx:
                label_list_aligned.append(label_ids[word_idx])  # 词的第一个subtoken
            else:
                label_list_aligned.append(-100)  # 词的后续subtoken忽略
            previous_word_idx = word_idx
        labels.append(label_list_aligned)
    tokenized["labels"] = labels
    return tokenized

# ---- 4. 基于词典的快速BioNER（规则方法）----
def dict_based_ner(text, gene_dict, disease_dict):
    """
    基于词典的生物医学实体识别（简单但快速）
    gene_dict: 基因名集合
    disease_dict: 疾病名集合
    """
    results = []
    words = text.split()
    for i, word in enumerate(words):
        word_clean = word.strip('.,;')
        if word_clean in gene_dict:
            results.append({'word': word_clean, 'type': 'Gene', 'pos': i})
        elif word_clean in disease_dict:
            results.append({'word': word_clean, 'type': 'Disease', 'pos': i})
    return results

# 示例词典（实际使用NCBI Gene、MeSH等专业数据库）
genes = {"BRCA1", "TP53", "EGFR", "KRAS", "HER2"}
diseases = {"cancer", "diabetes", "Alzheimer", "hypertension"}
text = "BRCA1 mutations are linked to breast cancer risk."
entities = dict_based_ner(text, genes, diseases)
for ent in entities:
    print(f"  [{ent['type']}] {ent['word']}")