Building a Healthcare Recommender with Keras, Two Towers and Google Cloud

Recommender systems are everywhere in entertainment and ecommerce: Netflix suggests your next show, Spotify your next song, Amazon your next purchase. The math behind those systems , taking a free text query and matching it against a big catalog , is exactly what you need to match ”I have shortness of breath, dizziness and palpitations” against a catalog of diseases.

That is the idea behind the Prescription Recommender App I want to walk through in this article. It takes plain language symptoms as input and returns the most likely diseases, along with suggested medications, diets and workouts. It is not a replacement for a doctor — it is a fast first step: ”here are the directions worth exploring.”

Under the hood it uses Keras Recommenders (the modern Kerasbased version of TensorFlow Recommenders), the TwoTower retrieval architecture, ScaNN for fast nearestneighbour search, a small Flask API served on Google Cloud Run, and it is designed to be deployed into Vertex AI Vector Search when it is time to scale. In this article I want to walk you through the whole thing, from raw CSV files to a production ready index.

The Value: what a Two Tower model actually buys you

The Two Towers Model

Before diving into the code, let me explain why this approach is a good fit for a domain like healthcare:

A traditional search engine matches keywords. If a user types ”my chest feels tight and I’m scared I’m going to die”, a keyword search will look for the literal words “chest”, “tight” and “die”. It will not understand that this is a textbook description of a panic attack. A Two Tower model, on the other hand, learns meaning. During training it places the phrase ”chest tightness, fear of dying” and the disease ”panic disorder” close to each other in a 768 dimensional space, even though they do not share a single word.

That is the first piece of value: semantic matching. The user does not need to know medical vocabulary. They describe what they feel, and the model handles the translation.

The second piece of value is speed. The Two Tower architecture is built specifically so that, after training, the candidate side (all the diseases) can be pre-computed once and stored in an index. At query time you only need to run the small query tower over the user’s text, and then do a nearest neighbor lookup.

In practice, a single query , from raw text to top 10 diseases, runs in tens of milliseconds on a single CPU. On Vertex AI Vector Search, with a properly sized endpoint, you can expect single digit millisecond latency even over millions of candidates. For a clinical support tool, that matters.

The third piece of value is scalability. The current dataset has a few hundred diseases, but the architecture does not care. Ten thousand diseases? Same code. A hundred thousand? Same code, just a bigger index. Because ScaNN (and Vertex AI Vector Search , which is built on the same algorithm) is an approximate nearest neighbor method, the search time grows roughly logarithmically with the size of the catalog instead of linearly. That is the difference between a demo and a product.

The fourth piece of value is separation of concerns. The model does one thing: it turns text into vectors that capture medical meaning. Everything else, the medications, the diets, the workouts, the explanations, live in regular dataframes that a domain expert can edit without touching any machine learning code. A medical reviewer can update a medication list in a CSV file, and the next deploy picks it up. No retraining needed. That is how you keep a healthcare product maintainable.

The Data: Six CSV files and a lot of cleaning

The dataset is a set of six CSV files assembled from public disease symptom datasets (https://www.kaggle.com/datasets/behzadhassan/sympscan-symptomps-to-disease). Here is what each file contains:

• description.csv : One short paragraph describing each disease
• Diseases_and_Symptoms_dataset.csv : A wide binary matrix: one row per disease, one column per symptom (1 = present)
• diseases_symptoms.csv : A second, freetext symptom list per disease
• medications.csv : A list of typical medications per disease
• diets.csv: Recommended foods and dietary advice
• workout.csv: Recommended physical activities

The first problem is that the files do not agree with each other. Some use ‘Disease’, others use ‘diseases’. Some have ‘Panic disorder’, others have ‘Panic Disorder’ with a capital D, and a few have trailing spaces. Anyone who has ever merged real world CSV files knows exactly what I mean.

The fix is boring but essential. Standardize the column names, lowercase everything, strip the whitespace:

import pandas as pd

df_desc = pd.read_csv('dataset/description.csv')
df_symp = pd.read_csv('dataset/Diseases_and_Symptoms_dataset.csv')
df_meds = pd.read_csv('dataset/medications.csv')
df_diets = pd.read_csv('dataset/diets.csv')
df_workouts = pd.read_csv('dataset/workout.csv')
df_additional_symptoms = pd.read_csv('dataset/diseases_symptoms.csv')

for d in [df_desc, df_symp, df_meds, df_diets, df_workouts]:
    col = 'diseases' if 'diseases' in d.columns else 'Disease'
    d.rename(columns={col: 'Disease'}, inplace=True)
    d['Disease'] = d['Disease'].astype(str).str.lower().str.strip()

The binary symptom matrix ( Diseases_and_Symptoms_dataset.csv ) needs a second transformation. Instead of keeping hundreds of 0/1 columns, we collapse them into a single text field listing the symptoms that are active for each disease, replacing underscores with spaces so the tokenizer can do its job:

symptom_cols = df_symp.columns[1:]

def get_active_symptoms(row):
    active = [col for col in symptom_cols if row[col] == 1]
    return ", ".join(active).replace('_', ' ')

df_symp['Symptoms_Text'] = df_symp.apply(get_active_symptoms, axis=1)

Then we merge the two symptom sources together (the binary derived one and the free text one), which gives each disease a rich, comma-separated symptom string. Finally we group everything by disease so that each row is unique:

df_candidates = (
    df_symp[['Disease', 'Symptoms_Text']]
    .merge(df_desc, on='Disease', how='left')
    .merge(df_additional_symptoms, on='Disease', how='left')
    .fillna("None")
)

df_grouped = df_candidates.groupby('Disease').agg({
    'Symptoms_Text': get_unique_items,
    'Description': get_unique_items,
}).reset_index()

The result is a clean dataframe where each disease has one description (the candidate for our recommender) and one long symptom text (the query we will train against). The last step is to wrap it in a ‘tf.data.Dataset’, which is what Keras expects:

import tensorflow as tf

dataset = tf.data.Dataset.from_tensor_slices(dict(df_unique_candidates)).map(lambda x: {
    "Disease": x["Disease_Description"],
    "Description": x["Symptoms_Text"],
})

shuffled = dataset.shuffle(df_grouped.shape[0], seed=42, reshuffle_each_iteration=False)
train = shuffled.take(int(df_grouped.shape[0] 0.95))
test = shuffled.skip(int(df_grouped.shape[0] 0.95))

A 95/5 split is fine here because the goal of training is not really to generalize to unseen diseases — all diseases will be in the index at inference time — but to learn the mapping from symptom language to disease meaning.

The Two Towers

The Two Tower architecture is two small neural networks that process two different types of input and produce vectors in the same shared space.

The query tower turns a free text symptom description into a 768-dimensional vector. It is built from three Keras layers: a TextVectorization layer that splits the text into tokens and maps them to integer IDs, an Embedding layer that turns those IDs into vectors, and a GlobalAveragePooling1D layer that averages the token vectors into a single sentence vector.

max_tokens = 5000
embedding_dimension = 768

vectorizer = tf.keras.layers.TextVectorization(
    max_tokens=max_tokens,
    output_mode="int",
    output_sequence_length=30,
)
vectorizer.adapt(train.map(lambda x: x["Description"]))

description_model = tf.keras.Sequential([
    vectorizer,
    tf.keras.layers.Embedding(max_tokens, embedding_dimension, mask_zero=True),
    tf.keras.layers.GlobalAveragePooling1D(),
])

Why 768 dimensions? It is the same width that BERT-base uses, and it gives a good balance between expressive power and index size. Smaller values (64, 128) train faster but cause more collisions between similar but distinct diseases. Larger values (1024) are overkill for a few hundred candidates.

The candidate tower is even simpler. Since each candidate is a full disease description (not a short ID), we use a StringLookup to turn it into an integer and an Embedding to give it a learnable vector:

import numpy as np

unique_diseases = np.unique(df_unique_candidates["Disease_Description"].values)

disease_model = tf.keras.Sequential([
    tf.keras.layers.StringLookup(vocabulary=unique_diseases, mask_token=None),
    tf.keras.layers.Embedding(len(unique_diseases) + 1, embedding_dimension),
])

Notice that the two towers do not share any weights. They are free to learn completely different things: the Query Tower becomes a symptom-language specialist, the Candidate Tower becomes a disease-name specialist. The only thing that ties them together is the training objective, which is the job of tensorflow recommenders.

Training: Teaching the Towers to Agree

The tfrs.tasks.Retrieval task does the following: for each pair in a batch, it computes the dot product between the query embedding and every candidate embedding in that batch, and pushes the model to make the correct pairs score higher than the incorrect ones. This is the same in batch negative sampling trick that powers large scale search systems at companies like YouTube and Spotify.

import tensorflow_recommenders as tfrs

class RecommenderModel(tfrs.Model):
    def __init__ (self, query_model, candidate_model):
        super(). __init__ ()
        self.query_model = query_model
        self.candidate_model = candidate_model
        self.task = tfrs.tasks.Retrieval(
            metrics=tfrs.metrics.FactorizedTopK(
                candidates=dataset.map(lambda x: x["Disease"]).batch(128).map(self.candidate_model)
            )
        )

    def compute_loss(self, features, training=False):
        query_embeddings = self.query_model(features["Description"])
        candidate_embeddings = self.candidate_model(features["Disease"])
        return self.task(query_embeddings, candidate_embeddings)

model = RecommenderModel(description_model, disease_model)
model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))

The FactorizedTopK metric tells you, during training, what fraction of the time the correct disease appears in the Top-1, Top-5, Top-10 and Top-50 results. Watching Top-10 accuracy climb from near zero to over 90% during training is one of the small joys of working with this library.

The script runs for 18,000 epochs. That sounds like a lot, and for a large dataset it would be absurd, but here each “epoch” is a single pass over a few hundred diseases, so the whole training run still finishes in a reasonable time on a modest GPU. You can use the Colab extension for VS Code with GPU for faster training.

cached_train = train.shuffle(df_unique_candidates.shape[0]).batch(1000).cache()
model.fit(cached_train, epochs=18000, verbose=0, callbacks=[tensorboard_callback, print_callback])

Building the ScaNN Index

Once the model is trained, the query tower is the only part we need at inference time. The candidate tower did its job: we can run it over every disease once, get a matrix of embeddings, and build a search index on top of that matrix.

The simplest option is to compare the user’s query vector against every disease vector one by one, an exact dot-product search. For a few hundred diseases that would actually be fine. But this pipeline is designed to be ready for hundreds of thousands of candidates, which is where ScaNN ( Scalable Nearest Neighbours , from Google Research) enters:

import numpy as np
import scann

symptoms_tensor = tf.constant(df_grouped["Symptoms_Text"].values)
embeddings = description_model(symptoms_tensor).numpy().astype(np.float32)

norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
normalized_dataset = embeddings / (norms + 1e9)

searcher = scann.scann_ops_pybind.builder(
    normalized_dataset, 10, "dot_product"
).tree(
    num_leaves=100,
    num_leaves_to_search=100,
    training_sample_size=len(normalized_dataset),
).score_ah(
    2, anisotropic_quantization_threshold=0.2,
).reorder(100).build()

First, we normalize every vector to unit length before building the index, what turns dot-product similarity into cosine similarity, which is what you almost always want for text embeddings.

Second, the tree().score_ah().reorder() chain is ScaNN’s classic recipe:

• partition the space into a tree
• score each partition with asymmetric hashing (a clever form of quantization)
• and then rerank the top candidates with exact dot products.

The result is search quality that is almost indistinguishable from brute force, at a fraction of the cost.

Finally, we save everything we need for inference:

import os

ARTIFACTS_DIR = "saved_artifacts"
os.makedirs(ARTIFACTS_DIR, exist_ok=True)

description_model.save(os.path.join(ARTIFACTS_DIR, "description_model"))
searcher.serialize(os.path.join(ARTIFACTS_DIR, "scann_index"))
df_grouped.to_pickle(os.path.join(ARTIFACTS_DIR, "df_grouped.pkl"))
df_meds.to_pickle(os.path.join(ARTIFACTS_DIR, "df_meds.pkl"))
df_diets.to_pickle(os.path.join(ARTIFACTS_DIR, "df_diets.pkl"))
df_workouts.to_pickle(os.path.join(ARTIFACTS_DIR, "df_workouts.pkl"))
np.save(os.path.join(ARTIFACTS_DIR, "normalized_embeddings.npy"), normalized_dataset)

Notice what is being saved:

• the query tower (as a Keras SavedModel)
• the ScaNN index
• the three lookup dataframes
• and the raw normalized embeddings

That last file is important. If you ever want to rebuild the index with different parameters — or ship it to Vertex AI Vector Search — you do not need to retrain anything. You just reload the .npy file.

Inference: the Flask endpoint on Cloud Run

At inference time, the app is very simple. Load the artifacts once at startup, then each incoming request goes through three steps: Embed, Search, Enrich.

import tensorflow as tf
import scann
import numpy as np
import pandas as pd

description_model = tf.keras.models.load_model("saved_artifacts/description_model")
searcher = scann.scann_ops_pybind.load_searcher("saved_artifacts/scann_index")
df_grouped = pd.read_pickle("saved_artifacts/df_grouped.pkl")

def embed_and_normalize(text: str) > np.ndarray:
    vec = description_model(tf.constant([text])).numpy().astype(np.float32)
    norm = np.linalg.norm(vec, axis=1, keepdims=True)
    return vec / (norm + 1e9)

@app.route("/search", methods=["POST"])
def search():
    query = request.get_json()["query"].strip()
    query_vec = embed_and_normalize(query)
    neighbors, scores = searcher.search(query_vec[0], final_num_neighbors=5)

    results = []
    for rank, (idx, score) in enumerate(zip(neighbors, scores), start=1):
        row = df_grouped.iloc[int(idx)]
        disease = row["Disease"].title()
        meds, diets, workouts = lookup_supplementary(row["Disease"])
        results.append({
            "rank": rank, "disease": disease, "score": round(float(score), 4),
            "medications": meds, "diets": diets, "workouts": workouts,
        })
    return jsonify(results)

The entire production inference path: embed one short piece of text, do a vector search, look up three dataframes, runs in tens of milliseconds on a single vCPU. Wrap it in a Dockerfile based on tensorflow/tensorflow:2.13.0, push it to Google Cloud Run , and you have an autoscaling endpoint without ever touching a Kubernetes cluster. The project’s actual Dockerfile does exactly that, using gunicorn as the WSGI server and exposing port 8080, which is what Cloud Run expects.

Cloud Run is the current deployment target for this project, and it is a great fit for a small to medium recommender: the container holds the query tower, the ScaNN index and the enrichment dataframes, and may scale to zero (beware of the cold start) when nobody is using it. For many use cases, that is all you will ever need.

The next step: Vertex AI Vector Search

Cloud Run works beautifully for a demo or a low traffic production app, but the moment the catalog grows, more diseases, multiple languages, millions of queries a day , the ScaNN index must live in a dedicated vector database, and on Google Cloud that database is Vertex AI Vector Search (formerly Matching Engine).

What is beautiful about this migration is that Vertex AI Vector Search is built on the same ScaNN algorithm we are already using locally. Moving from one to the other is not a rewrite, is a deployment change. The embeddings saved to normalized_embeddings.npy are the input to Vertex AI Vector Search. Here is a sketch of the migration:

from google.cloud import aiplatform
import json

aiplatform.init(project="my-project", location="us-central1")

# 1. Upload the embeddings as JSONL to Google Cloud Storage
with open("embeddings.jsonl", "w") as f:
    for i, vec in enumerate(normalized_dataset):
        f.write(json.dumps({"id": str(i), "embedding": vec.tolist()}) + "\n")

# 2. Create the index
index = aiplatform.MatchingEngineIndex.create_tree_ah_index(
    display_name="prescriptionrecommender",
    contents_delta_uri="gs://mybucket/embeddings/",
    dimensions=768,
    approximate_neighbors_count=10,
    distance_measure_type="DOT_PRODUCT_DISTANCE",
)

# 3. Deploy it to an endpoint
endpoint = aiplatform.MatchingEngineIndexEndpoint.create(
    display_name="prescriptionendpoint",
    public_endpoint_enabled=True,
)
endpoint.deploy_index(index=index, deployed_index_id="prescription_v1")

Once the index is deployed, the Flask app no longer needs to load ScaNN at all. It just keeps the Query Tower (a few megabytes) and calls the Vertex AI endpoint:

response = endpoint.find_neighbors(
    deployed_index_id="prescription_v1",
    queries=[query_vec[0].tolist()],
    num_neighbors=5,
)

The index is now separately scalable and you can add nodes without redeploying the app. It is updatable in place: adding a new disease is a matter of appending one embedding, not rebuilding a container. It is multiregion if you need it to be. And because the query tower runs on CPU and is extremely small, you can keep serving it from Cloud Run while the heavy vector lookup happens on a managed, dedicated service. We have stateless inference code on Cloud Run and stateful vector index on Vertex AI Vector Search.

You can start with a single Python script and a Flask app on your laptop, and the same artifacts, the same embeddings, the same 768 dimensional space, the same distance metric, carry you all the way to a production, multiregion, millisecond latency vector search service with just a redeploy.

Closing Remarks

This project combines three things that are usually kept in separate boxes:

• classic machine learning (the TwoTower architecture has been around for years, I started using it in 2022)
• modern Keras (the code above runs on Keras and TensorFlow 2.13)
• and Google Cloud’s managed infrastructure, Cloud Run and Vertex AI Vector Search.

If you are a developer thinking about how to bring semantic search into a domain you care about , healthcare, legal, education, customer support, the gap between ‘interesting tutorial’ and ‘something you could actually put in front of users’ is much smaller than it looks. You need a clean dataset (essential), two small Keras models, a ScaNN index, and a couple of hundred lines of Python.

Run the Python script and you will have your own symptom to disease recommender running locally in minutes. Build the Docker image, push it to Cloud Run, and it is online. When the catalog grows, point the same embeddings at Vertex AI Vector Search, and it is ready for scale.

That is the whole idea: start simple, stay in the same ecosystem, and let the platform grow with the product.

Acknowledgements

✨ Google ML Developer Programs and Google Developers Program supported this work by providing Google Cloud Credits (and awesome tutorials for the Google Developer Experts)✨

🔗https://developers.google.com/machinelearning 🔗