BioNERBERT-CRF Entity Recongition
Created an inference pipeline for a custom Biomedical Named Entity Recognition (NER) model with a Conditional Random Fields (CRF) layer. The model extracts key entities from unstructured text. The model is fine-tuned for the following entities:
- CELL_FUNCTION
- CELL_TYPE
- CONDITION
- LOCATION
- MOLECULE
- PATHOGEN
- PROCESS
- RECEPTOR
- TREATMENT
The pipeline was previously hosted on Google Cloud Run which would automatically scale to zero It is now hosted on a Microsoft Azure B2als v2 virtual machine with automated shut down and start up procedures to minimize infrastructure costs.
This model was created with flawed training data annotated by GPT-4 mini, and the data that it was trained on is very lacking. This was a project meant to fine-tune a NER model with, so it definitely should not be used for anything important. See the Hugging Face page for more details.
1. MLOps Architecture and Serverless Constraints
Originally, this pipeline was hosted on Google Cloud Run, I then migrated this over to a Microsoft Azure virtual machine to share the server costs with the Timmy and Tommy project, which did not need the full 4 GB of RAM that the virtual machine came with. Standard managed serverless inference endpoints (e.g., Hugging Face Inference Endpoints) enforce rigid model architecture constraints for security and sandboxing. Because this custom NER model utilizes a Conditional Random Fields (CRF) layer stacked upon the fine-tuned PubMedBERT backbone, standard free-tier inference APIs flag the execution of the custom CRF logic.
To bypass this architectural limitation without spinning up expensive persistent EC2/GCE instances, a custom serving layer was engineered using FastAPI and PyTorch. The entire runtime was containerized with Docker and deployed to a Microsoft Azure Virtual Machine allowing arbitrary model architectures to execute in a serverless, horizontally scaling environment.
2. Scalability and Cold Starts
To strictly minimize costs, the Cloud Run instance scales to exactly 0 when inactive. Upon receiving
a cold-start invocation, the container initializes the FastAPI server and pulls the ~400MB
.pt weights directly from the Hugging Face Hub at runtime before beginning inference.
This incurs a deliberate latency penalty of approximately 60 seconds during a cold start, trading
execution speed for absolute cost efficiency.
sequenceDiagram
participant User
participant Server as Cloud Run (FastAPI + PyTorch)
participant HF as Hugging Face Hub
User->>Server: Enter website
Server-->>User: Render UI
User->>Server: Submit Text
alt If Server was off for a while (Cold Start)
Server->>HF: Download weights from hugging face
end
Server->>Server: NER model analyzes text
Server-->>User: 6. Display Results
3. Model Details
Rather than memorizing all of the possible treatments, cell types, molecules, etc. the model performs
its identification by learning how humans talk about each of these items. By doing this instead of
raw memorization, the model will be able to identify new entities that weren't in its training data,
as long as how people talk about these things doesn't change.
Below are two images of a sentence with real drugs and conditions, and one with completely made up
names (definitely not in its training data).The BioNERBERT model is able to identify everything as
expected. .
To confirm that these aren't real medical terms that the model has memorized, I confirmed with Chat
GPT. To be one hundred percent sure yourself, you can also try the site with the link at the top of
the page.
While there are benefits with making the model understand the structure rather than memorize
everything, there are also some issues with the current pipeline. I used ChatGPT-mini to annotate
all of these
entities in text that I scraped from PubMedCentral (the PMC database), as a result, many of its
flaws also made it into the BioNERBERT pipeline. "Process" was labelled very infrequently,
resulting in the fine-tuned model being very conservative in predicting treatment. Additionally,
there is another entity that had a single label named "VISUAL_PROPERTY", which is practically
impossible to trigger.