from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Load the fine-tuned model and tokenizer
model = AutoModelForSequenceClassification.from_pretrained("your_model_path")
tokenizer = AutoTokenizer.from_pretrained("dmis-lab/biobert-v1.1")

def predict(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=256)
    with torch.no_grad():
        logits = model(**inputs).logits
    return "urgent" if torch.argmax(logits) == 1 else "non-urgent"

# Example usage
print(predict("Patient has chest pain and difficulty breathing."))

Training Details

Training Data

Dataset: myothiha/triage_dataset

Training Procedure

Training Hardware: 2 GPUs
Batch Size: 48
Learning Rate: 6e-6
Epochs: 6
Loss Function: Weighted cross-entropy (to handle class imbalance)

Training Hyperparameters

Optimizer: AdamW
Gradient Accumulation Steps: 3
Weight Decay: 0.05
Mixed Precision: Enabled (fp16)

Evaluation

Epoch	Training Loss	Validation Loss	Accuracy	Precision	Recall	F1 Score
1	No log	0.606408	0.681326	0.436587	0.656415	0.524395
2	0.639200	0.564773	0.706726	0.467656	0.692443	0.558271
3	0.639200	0.562810	0.718721	0.481095	0.648506	0.552395
4	0.551200	0.559481	0.701317	0.462286	0.710896	0.560249
5	0.522800	0.560782	0.710019	0.471334	0.686292	0.558855
6	0.522800	0.559855	0.705786	0.466978	0.702109	0.560899

Summary

This fine-tuned BioBERT model helps classify medical questions by urgency. It can be used in chatbots, triage systems, or AI-powered healthcare assistants to improve response prioritization.

Feel free to test it and integrate it into your applications!

Yuvrajxms09/biobert-triage-classifier

BioBERT Triage Classifier

Model Description

Model Sources

Uses

Direct Use

Out-of-Scope Use

Recommendations

How to Get Started with the Model