There’s a specific kind of frustration that comes from having a model that works perfectly in a notebook and then spending three days figuring out how to actually deploy it. That was my starting point with MedGemma.
MedGemma is Google’s medical-domain multimodal model – built for clinical text, patient notes, that sort of thing. It’s not a simple classification model you can wrap in two lines and call it done. The inference logic is opinionated, the inputs need careful handling, and getting consistent JSON out of it required a few tricks I had to figure out the hard way.
Here’s how I got it packaged, registered, and deployed using MLflow’s PyFunc framework on Databricks.
First, Getting the Model Down Locally
MedGemma is a gated model on Hugging Face, so before anything else you need a token with access approved. Once that’s sorted, I downloaded the whole thing using snapshot_download into a temp directory on the cluster.
python
from huggingface_hub import snapshot_download
import os, shutil
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
LOCAL_TMP_PATH = "/tmp/medgemma_download"
if os.path.exists(LOCAL_TMP_PATH):
shutil.rmtree(LOCAL_TMP_PATH)
os.makedirs(LOCAL_TMP_PATH, exist_ok=True)
snapshot_download(
repo_id="google/medgemma-4b-it",
local_dir=LOCAL_TMP_PATH,
token=os.environ["HUGGINGFACEHUB_API_TOKEN"],
local_dir_use_symlinks=False,
resume_download=True
)
The resume_download=True saved me more than once – the model is large enough that a network hiccup mid-download would otherwise mean starting over completely.
One thing worth noting: I installed the dependencies directly on the cluster rather than through a requirements file, since it was a dedicated cluster and I didn’t want to affect anything else running elsewhere. torch, transformers, accelerate, bitsandbytes, sentencepiece, mlflow – all pip installed at the top of the notebook.
Building the PyFunc Wrapper
This is the part that took the most thought. PyFunc gives you a clean interface – implement load_context and predict, and MLflow handles the rest. But what you put inside those methods is entirely up to you, which means all the complexity of the model lives in your wrapper class.
My wrapper is called MedGemmaModel and it extends mlflow.pyfunc.PythonModel.
Loading the model
load_context runs once when the model is first loaded, not on every prediction. That’s where I initialize the processor and model:
python
def load_context(self, context):
self.device = "cuda" if torch.cuda.is_available() else "cpu"
model_path = context.artifacts["model_path"]
self.processor = AutoProcessor.from_pretrained(model_path)
self.model = AutoModelForImageTextToText.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
device_map="auto",
low_cpu_mem_usage=True,
)
self.model.eval()
bfloat16 keeps memory usage reasonable without a meaningful quality hit for this use case. device_map=”auto” lets the framework figure out GPU placement rather than hardcoding it, which matters when you don’t always know exactly what hardware you’ll land on.
I also set token limits here rather than hardcoding them scattered through the prediction logic – max_input_tokens, chunk_size, overlap, max_new_tokens. Having them in one place made tuning much less painful
The Chunking Problem
Clinical notes are long. Sometimes very long. A standard transformer context window handles a certain amount of text cleanly, but once you go past that things degrade or break entirely.
My solution was to count tokens first, then decide how to handle the input:
python
token_ids = self.processor.tokenizer(
text, return_tensors="pt", truncation=False
)["input_ids"][0]
if len(token_ids) <= self.max_input_tokens:
# single pass
else:
# chunk it
For long inputs, I split the token ID tensor into overlapping chunks – 6000 tokens each with a 300-token overlap so context doesn’t get hard-cut at chunk boundaries. Each chunk gets processed independently, producing its own partial JSON output. Then everything gets merged at the end.
The merge logic handles conflicts sensibly: equal values stay as-is, two lists get unioned without duplicates, a list and a scalar either append or promote to a list if there’s a conflict. It’s not perfect but it covers the realistic cases well.
Forcing JSON Output
This was the trickiest part. MedGemma, like most LLMs, doesn’t always open with { even when you ask it to return JSON. It might add a preamble, wrap things in markdown fences, or just narrate what it’s about to do before doing it.
I handled this by literally appending { to the input token sequence before generation:
python
brace_ids = self.processor.tokenizer("{", add_special_tokens=False)["input_ids"] brace_tensor = torch.tensor([brace_ids], device=device)
inputs["input_ids"] = torch.cat([inputs["input_ids"], brace_tensor], dim=-1)
Then after generation, I prepend that { back onto the decoded output and run it through a bracket-balanced JSON extractor rather than trying to regex my way through potentially nested structures. The extractor walks character by character, tracks depth, handles strings and escape sequences, and pulls out the first complete JSON block it finds.
Does it always work? Mostly. Edge cases exist. But it’s dramatically more reliable than hoping the model behaves.
Logging and Registering the Model
Once the wrapper was working, packaging it for MLflow was straightforward:
python
with mlflow.start_run() as run:
mlflow.pyfunc.log_model(
artifact_path=model_name,
python_model=MedGemmaModel(),
artifacts={"model_path": LOCAL_TMP_PATH},
signature=signature,
pip_requirements=[
"torch>=2.1.0",
"transformers>=4.50.0",
"accelerate",
"bitsandbytes",
"sentencepiece",
"pandas"
]
)
run_id = run.info.run_id
mlflow.register_model(
model_uri=f"runs:/{run_id}/{model_name}",
name=“<catalog>.huggingface_ml.medgemma_4b"
)
The artifacts dictionary is important here. Rather than downloading the model again at serve time – which would be slow and potentially unreliable – I point MLflow at the local path where the model already lives. MLflow packages that path as part of the logged artifact, and context.artifacts[“model_path”] gives it back to you cleanly inside load_context.
The signature was inferred from a simple input/output example – a DataFrame with input and optionally prompt columns going in, a DataFrame with a response column coming out.
Deploying via Databricks UI
For serving, I used the Databricks Model Serving UI rather than scripting the deployment. Mostly because I wanted visibility into what was happening – the build logs, the environment setup, whether the GPU was being picked up correctly.
Step-by-Step: Deploying via Databricks UI
- Confirm the model is registered After your mlflow.register_model() call runs successfully, go to Models in the left sidebar and confirm raw.huggingface_ml.medgemma_4b shows up under Unity Catalog, with the version number you just registered.
- Open the Serving tab Click Serving in the sidebar, then click Create serving endpoint.
- Name the endpoint Give it a descriptive name (e.g. medgemma-4b-endpoint). One thing to know: endpoint names can’t use the databricks-prefix — that’s reserved for Databricks’ own preconfigured endpoints.
- Select the served entity
Click into the Entity field and choose My models – Unity Catalog (since your model is registered there, not the legacy Workspace Model Registry). Then pick:
- Catalog: raw
- Schema: huggingface_ml
- Model: medgemma_4b
- Version: whichever version you just logged
- Choose compute Since MedGemma needs a GPU, select a GPU compute size rather than CPU. Pick based on what your model actually needs — for a 4B parameter model in bfloat16, you don’t need the largest tier, but check your memory footprint before deciding.
- Configure env. Variable if needed
- Configure scale-out (optional) Set minimum/maximum provisioned concurrency if you want autoscaling behavior, or leave it at defaults for a first deployment.
- Create the endpoint
Note: The build logs were genuinely useful here. The first attempt failed because of a dependency conflict I hadn’t caught – visible immediately in the logs rather than something I’d have had to dig out of a CLI response. Watching the container build step by step also gave me confidence that the GPU configuration was being applied correctly, which with a model this size isn’t something you want to discover is wrong after deployment.
Once the endpoint was live, testing it was a simple POST request with a DataFrame payload.
What PyFunc Actually Buys You Here
Looking back at this whole thing – the chunking, the JSON forcing, the merge logic – none of it fits neatly into a standard model-serving pattern. There’s no sklearn wrapper that handles overlapping token chunks. There’s no built-in serving format for “run this logic, call this model, parse the output, merge across chunks.”
PyFunc let me write all of that as regular Python and still end up with something that Databricks treats like any other registered model. Same deployment flow, same endpoint format, same versioning. That’s the actual value – not that it made any individual piece easier, but that it meant I didn’t have to build a separate deployment system just because my inference logic was complicated.
