fhir4px-embeddings-onnx

fhir4px Embeddings ONNX

Pre-converted ONNX versions of embedding models used by fhir4px for Tier 3 categorization.

Models

pubmedbert-base-embeddings/

NeuML/pubmedbert-base-embeddings converted to ONNX format for use with transformers.js.

Files:

• onnx/model.onnx — fp32 (416MB)
• onnx/model_quantized.onnx — int8 dynamic quantization (105MB, production)
• config.json, tokenizer.json, tokenizer_config.json, vocab.txt, special_tokens_map.json

Why this model:

• Fine-tuned from Microsoft PubMedBERT for sentence-similarity on PubMed title-abstract pairs
• Apache 2.0 license
• Best accuracy on fhir4px categorization tasks (see docs/INTEGRATION.md)
• Smaller than alternatives at int8 (105MB vs 143MB for gte-modernbert-base q8)

Usage:

import { pipeline, env } from "@huggingface/transformers";
env.allowRemoteModels = true;
const extractor = await pipeline(
  "feature-extraction",
  "joelmontavon/fhir4px-embeddings-onnx",
  {
    dtype: "q8",
    subfolder: "pubmedbert-base-embeddings",
  }
);
const output = await extractor(texts, { pooling: "mean", normalize: true });

Precomputed Centroids (centroids/)

Each classification task has a JSON file with precomputed centroids (768-dim, L2-normalized). When the centroid field is present, the runtime classifier skips prototype embedding entirely.

Files:

• centroids/observation_category.json — lab / vital / other (34 prototypes)
• centroids/allergy_type.json — medication / food / environmental / other (32 prototypes)
• centroids/visit_type.json — inpatient / outpatient / emergency / telehealth / procedure (25 prototypes)

Format:

{
  "task": "observation_category",
  "model": "joelmontavon/fhir4px-embeddings-onnx",
  "classes": {
    "lab": {
      "centroid": [0.0123, -0.0456, ...],
      "prototype_texts": ["Hemoglobin A1c", "Glucose", ...]
    },
    ...
  }
}

How computed:

• Model: onnx/model_quantized.onnx (q8 dtype, matching runtime inference path)
• Pooling: mean pooling with attention mask
• Normalization: L2 normalized embeddings, then L2 normalized centroid (mean of prototype embeddings)
• All centroids are unit vectors (L2 norm = 1.000000)

Usage in browser runtime:

const response = await fetch(
  `https://huggingface.co/${MODEL_REPO}/resolve/main/centroids/observation_category.json`
);
const { classes } = await response.json();

// For each class, use classes[className].centroid directly
// (skip embedding prototype_texts since centroid is precomputed)

Regenerating: run python scripts/embedding_eval/compute_centroids.py after updating prototype texts.