medBERT-base

This repository contains a BERT-based model, medBERT-base, fine-tuned on the gayanin/pubmed-gastro-maskfilling dataset for the task of Masked Language Modeling (MLM). The model is trained to predict masked tokens in medical and gastroenterological texts. The goal of this project is to improve the model's understanding and generation of medical-related information in natural language contexts.

Model Architecture

• Base Model: bert-base-uncased
• Task: Masked Language Modeling (MLM) for medical texts
• Tokenizer: BERT's WordPiece tokenizer

Usage

Loading the Pre-trained Model

You can load the pre-trained medBERT-base model using the Hugging Face transformers library:

from transformers import BertTokenizer, BertForMaskedLM
import torch

tokenizer = BertTokenizer.from_pretrained('suayptalha/medBERT-base')
model = BertForMaskedLM.from_pretrained('suayptalha/medBERT-base').to("cuda")

input_text = "Response to neoadjuvant chemotherapy best predicts survival [MASK] curative resection of gastric cancer."
inputs = tokenizer(input_text, return_tensors='pt').to("cuda")

outputs = model(**inputs)

masked_index = (inputs['input_ids'][0] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0].item()

top_k = 5
logits = outputs.logits[0, masked_index]
top_k_ids = torch.topk(logits, k=top_k).indices.tolist()
top_k_tokens = tokenizer.convert_ids_to_tokens(top_k_ids)

print("Top 5 prediction:")
for i, token in enumerate(top_k_tokens):
    print(f"{i + 1}: {token}")

Top 5 prediction: 1: from 2: of 3: after 4: by 5: through

Fine-tuning the Model

To fine-tune the medBERT-base model on your own medical dataset, follow these steps:

1. Prepare your dataset (e.g., medical texts or gastroenterology-related information) in text format.
2. Tokenize the dataset and apply masking.
3. Train the model using the provided training loop.

Here's the training code:

https://github.com/suayptalha/medBERT-base/blob/main/medBERT-base.ipynb

Training Details

Hyperparameters

• Batch Size: 16
• Learning Rate: 5e-5
• Number of Epochs: 1
• Max Sequence Length: 512 tokens

Dataset

• Dataset Name: gayanin/pubmed-gastro-maskfilling
• Task: Masked Language Modeling (MLM) on medical texts

Acknowledgements

• The gayanin/pubmed-gastro-maskfilling dataset is available on the Hugging Face dataset hub and provides a rich collection of medical and gastroenterology-related information for training.
• This model uses the Hugging Face transformers library, which is a state-of-the-art library for NLP models

Support: