chrisvoncsefalvay/CRAG-dual-encoder-mimicause

CRAG: Causal Reasoning for Adversomics Graphs

CRAG (Causal Reasoning for Adversomics Graphs) is a dual-encoder model for extracting adverse drug event (ADE) relationships from clinical narratives.

Model Description

CRAG uses a dual-encoder architecture with:

• Two separate BioLinkBERT encoders: One for drug mentions, one for adverse event mentions
• mean pooling: Sequence representation aggregation
• Bilinear fusion: Captures complex drug-ADR interactions
• Multi-view concatenation: Combines bilinear output, individual embeddings, and element-wise products

Training Approach

The model is trained in two phases:

1. Phase 1 - Contrastive Pre-training: InfoNCE loss with hard negative mining to learn discriminative drug-ADR embeddings
2. Phase 2 - Classification Fine-tuning: Focal loss to handle class imbalance and refine the classifier

Performance

Test Set Results

Metric	Score
F1 Score	0.9264705882352942
Precision	0.9163636363636364
Recall	0.9368029739776952
AUC-ROC	0.9656523724050712

Usage

from transformers import AutoTokenizer

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("michiyasunaga/BioLinkBERT-base")
tokenizer.add_special_tokens({"additional_special_tokens": ["[DRUG]", "[/DRUG]", "[ADR]", "[/ADR]"]})

# Download and load the model
from huggingface_hub import hf_hub_download
import torch

model_path = hf_hub_download(repo_id="chrisvoncsefalvay/CRAG-dual-encoder-mimicause", filename="pytorch_model.pt")
state_dict = torch.load(model_path, map_location="cpu")

Training Configuration

Parameter	Phase 1 (Contrastive)	Phase 2 (Classification)
Epochs	5	8
Batch Size	64	64
Learning Rate	2e-5	2e-5
Loss Function	InfoNCE (tau=0.07)	Focal (gamma=1.0, alpha=0.75)
Hard Negatives	50%	-

License

Apache 2.0