BioNERBERT-CRF Entity Recongition

Status
Loading telemetry...
Live
Live App ↗ (Available Weekdays from 08:00 - 18:00 EST)
Tech Stack FastAPI, PyTorch, Docker, Google Cloud Run, Microsoft Azure, Virtual Machines

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.

Model Details

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.

BioNERBERT Frontend

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. .

Made-Up Entities

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.