Why we think LLMs can be useful and why we will not replace all of our models with them
Context
At Medium, we have many Machine Learning models that we use to label stories automatically. These affect what stories we recommend to readers.
Here’s some examples:
Some Clarifications on our Machine Learning policy
Before we go deep on this project, I just wanted to clarify a few things about how we stand regarding AI in general.
Medium has been training internal models with user and post data for a long time now. We train models with specific tasks. For example, models that power our recommendations algorithm, or text classification models like the ones presented in this story. All in the goal to improve our product. With the LLM approach I describe in this story, we ARE NOT sharing these models with other companies. And we ARE NOT allowing anyone to train on our users’ data and content. Here we used Snowflake LLM tools for inference only (no LLM training was done here) and they are actually hosting all of the models inside their own infrastructure and guarantee that they are not using any of this for training. Shoutout to the Snowflake team for making it so easy and safe to use LLMs on our data!
If you want to read more about Medium’s stance on AI, I definitely recommend giving these a read:
- Default No to AI Training on Your Stories
- Finally, an internet standard for writers’ rights vs. AI companies
- We want your feedback: How can writers use AI to tell human stories?
Problem
During our roadmap planning we decided that our NSFW model was out of date and it was time to revamp it. This model labels stories as “Not Safe for Work” if they have sexually explicit content, lots of profanity, or basically anything you wouldn’t want to read on your big monitor in the middle of an open space!
As you can imagine it’s a pretty important model. We really need it to make sure our most “interesting” content only reaches our most “interested” users and ONLY them!
There’s actually a funny anecdote from 2021 that was shared internally back then. One of our employee was onboarding the WHITE HOUSE staff onto Medium so that they could start using the platform with the POTUS account. And the first thing that showed up on the homepage was a big picture of well… a butt… So that became of whole thing “Our algorithm is serving erotica to President Joe Biden!!”. That was a brand new account so that was our top story “by default”.
That’s embarrassing for us and for our users! This anecdote is actually what prompted the recommendations team to create the first version of our NSFW model (so, thank you Joe!).
We’re now in 2026 and this model is now pretty old. It’s not getting retrained regularly, the code runs on python versions that are extremely old. With turnover and changes in tools and file organization, we lost track of how it was trained and what kind of performance was measured at the time. So this is now an obscure part of the ML stack, in need of a good makeover.
LLM based approach
The typical way we’d build a binary classification model like this is to build a pipeline where:
- new data is regularly added to a training set by a human labeling team
- and the model is regularly retrained
So now, what if we replace all that with a really simple LLM based approach?
With LLMs it’s pretty easy to create whatever text classifier you want. You just need the proper prompt and have some sort of expected output that you can parse from the LLM’s response. And there’s no need for training at all.
We thought this NSFW model revamp was a great opportunity to give this a try. Are we able to build a NSFW classifier this way, that matches our performance expectations?
How we built it
The process was pretty simple
First we built an evaluation dataset:
- we sampled a random sample of stories published on Medium
- then we had our curation team label the dataset (is it NSFW? yes/no)
- as they did that, we also asked them to refine their definition of “NSFW”, and give concrete examples (this was very useful to build the prompt)
Side note: when we picked the random sampling of stories, we used other signals to make sure we oversample NSFW stories. If we just sample at random, we’d have to sample thousands of stories just to get a handful of NSFW examples. Thanks to this, our evaluation had a nice 50/50 split on NSFW/SFW labels.
Next step was to build the prompt for the LLM:
- we leveraged all the clarifications that the curators added while labeling the dataset and summarized that into “NSFW guidelines”
- we pass in the story’s contents
- we ask the LLM to output a score from 0 → 100
- and we also ask the LLM to give a short explanation for the score. For example: “graphic description of sexual intercourse”
Finally we need to evaluate the LLM approach and compare the different model:
- have a few LLMs make their predictions on the evaluation set
- have the legacy model make it’s predictions on the evaluation set too
- and then compare the metrics
Great news, after some testing, we found some LLM models that match or outperform the legacy model! We got parity with the legacy model on both false positives and false negatives. Some of the bigger LLMs even outperform the legacy model on both metrics. The mistral models performed really well for us, as well as some of the larger Claude models.
Performance is key here because there’s big downsides with both false negatives and false positives:
Now, we need to get an idea of the costs: is that LLM approach going to cost us an arm and leg?
Something nice with LLMs is that the costs are easy to estimate. Each model has a cost per input token and a cost per output token. Using that and a little bit of back of the enveloppe math you can quickly get a really good cost estimate.
We then picked the model that was the best fit for us in terms of costs and performance. For us that was Mixtral 8×7b (I promise that was a fair trial – nothing to do with the fact that I’m French!). The main way we kept the costs under control is by picking a relatively cheap model, and also by only scoring the stories that are eligible to get distributed in the first place (for example, no need to waste time on stories that have been flagged as spam).
For offline experimentation and for production classification, we simply used Snowflake. All of our data ends up in Snowflake and they make it super simple to run LLM inference directly inside a SQL query. There’s a complete function available that lets you call an LLM and pick your model as well as the expected response format.
Here’s the query I used during the evaluation process as an example:
set modelName = 'claude-haiku-4-5';
set promptName = 'raph_test';
set promptVersion = '1.3';
set responseFormat = '{
"type": "json",
"schema": {
"type": "object",
"properties": {
"score": {"type": "integer"},
"reason": {"type": "string"}
},
"required": ["score", "reason"]
}
}';
with raw_llm_responses AS (SELECT post_id,
title,
text,
SNOWFLAKE.CORTEX.COMPLETE(
$modelName,
-- build prompt: instructions and post information
ARRAY_CONSTRUCT(
OBJECT_CONSTRUCT(
'role',
'system',
'content',
prompts.prompt
),
OBJECT_CONSTRUCT(
'role',
'user',
'content',
'<title>' || COALESCE(title, 'N/A') || '</title>' || '\n' ||
'<text>' || COALESCE(text, 'N/A') || '</text>'
)
),
OBJECT_CONSTRUCT(
'temperature', 0,
'response_format', PARSE_JSON($responseFormat)
)
) AS llm_raw_response
FROM posts_to_evaluate
join medium.ml.nsfw_classifier_prompt as prompts
where prompts.name = $promptName
and prompts.version = $promptVersion),
-- parse results from the returned JSON
llm_responses_parsed AS (SELECT post_id,
title,
text,
llm_raw_response,
llm_raw_response:structured_output[0]:raw_message AS parsed_json,
TRY_CAST(parsed_json:score::STRING AS INT) AS score,
parsed_json:reason::STRING AS score_reason,
parsed_json::STRING AS response_as_string
FROM raw_llm_responses)
SELECT post_id,
score AS prediction,
score_reason,
response_as_string as raw_response,
$modelName AS model_name,
$promptName AS prompt_name,
$promptVersion AS prompt_version
FROM llm_responses_parsed;
In our tech stack, it’s really easy to setup a simple job that runs a Snowflake query on a schedule. We just leveraged that and we were able to successfully ship that model very quickly.
And now repeat?
Once that new model was up and running, we were really interested in applying that same formula to our other models. We first sat down to weigh the pros and cons of that new LLM based approach, here’s how we broke it down:
Generally the LLM approach is:
- easier to setup
- more expensive
- less adaptive
One thing to keep in mind too is that LLMs are “stuck at a point in time”. They have been trained at a certain date and they don’t know anything about what happened after that date. So they will not be adapting to new trends and performance is likely to drop over time. Depending on the use case, you might want to upgrade your model every now and then.
The NSFW model was the perfect use case:
- we don’t have a proper training set for NSFW content. And it would be costly to create and maintain one
- we don’t feel like we need to be super reactive to new trends. I don’t really have data to support that, this is more a judgment call: we think that NSFW stories 6 months from now will be similar to the ones that are published today.
- costs are low as long as we only score a portion of the stories published on Medium. Here we just need to score the stories that are eligible for distribution in the first place
All in all, that model revamp was a great success:
- we have a brand new NSFW model
- it performs really well
- we don’t need to maintain a training set
- it’s cheap to run
- and it’s really simple in terms of engineering
What about our other models then?
It turns out that the NSFW model is maybe our only model that’s a good use case for an LLM based approach:
- for our topic model, we get our dataset for free (we just train on the topics that users are adding to their stories). This means that we automatically adapt to new trends, which is really nice. Switching to an LLM based approach would feel like a downgrade
- for the spam model, we also want to adapt to new trends quickly. Spammers are always trying new things and they are always trying to trick the platforms to get visibility. So in this case we are willing to pay the cost of maintaining a training set, augmented with new data every day. Also since we want to score ALL the stories posted on Medium, it becomes more complex and costly to use an LLM based approach
What about new use cases?
Even though that doesn’t work really well for our other existing models. It is still very exciting because it makes it really easy to spin up new text classification models. What if we wanted to give users more control on their content, like “allow erotica” but do not allow “true crime”? We’re a text-based platform, so the possibilities are endless!
Some side notes
Why ask for scores and not just a yes/no value?
With scores, you are able to rank the predictions and it means that you can use metrics that quantify how well you separate the NSFW stuff from the SFW stuff (metrics like ROC-AUC or PR-AUC). It also means that you can set the threshold wherever you like and you can change that in the future too – letting you choose the balance of false positives and false negatives that suits you best.
Why ask for a short description?
Our end goal is just to know if a story is SFW or not, we don’t actually need the LLM to explain it’s decision. And it does add a little to the inference costs (the output becomes longer). But in our use-case it was negligible in the overall costs.
I had to do a lot of spot checking and debugging on individual stories for this project. (I learned a lot of things 😳). And having the descriptions alongside the score was super useful for that. It’s also a great thing to have now that this is production, and it can also be used to break down NSFW stories into categories. That can be useful for the curation team or the trust and safety team
To wrap this up, it was really fun to build this model and we’re excited to build new features with this new tool in our toolbox! We’re hoping this opens up new possibilities for where our product teams can take Medium next and helps us make sure readers find more stories they love, and writers get more readers that love their work!
How we think about text classification in the LLM era was originally published in Medium Engineering on Medium, where people are continuing the conversation by highlighting and responding to this story.
Source: medium.engineering
