How we think about text classification in the LLM era

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:

Three examples of models that classify stories based on their content. An NSFW model and a Spam model each produce a single score, and a threshold decides the label — NSFW or SFW, Spam or Ham. A Topics model instead scores a long list of topics at once, and the ones scoring high enough are selected as the story’s topics. — a few of our text classification models. All diagrams and charts made by the author

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:

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

The typical machine learning setup, in two loops. In the training loop, humans label new data, those labels get added to the training set, the model is retrained on a daily or weekly cadence, and the retrained model is published. In the prediction loop, new data arrives — a story published on Medium — the latest published model makes a prediction on it, and the prediction is stored.

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.

The LLM-based approach. Four pieces are assembled into a single prompt: instructions asking whether the post qualifies as NSFW and requesting a 0-to-100 score with an explanation, a set of NSFW guidelines, a handful of labelled examples, and the post’s content. The LLM returns a score and a written explanation. For example:score=12 expalanation=“nothing matches NSFW guidelines” — We build a prompt with all the context needed for the LLM to make a decision. We get back a score between 0 (perfectly SFW) and 100 (undoubtedly NSFW) as well as a short explanation of the score.

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

How the evaluation works. A dataset of story contents with a human-assigned yes/no NSFW label is run through each candidate model — the legacy model, Claude Haiku, Mixtral 8x7b and others — and each one gets a precision and recall score that can be compared side by side. Most land around 95% on both, while the smaller Claude Haiku shows high precision but much lower recall. — Thanks to our new evaluation dataset, we can compare our legacy model with LLM models

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:

A confusion matrix showing why both kinds of error hurt. When a story is actually SFW but the model flags it NSFW, that’s a false positive: the story doesn’t get distributed even though it should have, and the writer loses out. When a story is actually NSFW but the model flags it SFW, that’s a false negative: it might get distributed to users who don’t expect it — the Joe Biden problem. The other two quadrants, true positives and true negatives, are the correct outcomes.

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.

A three-step simplification of the LLM cost equation. The full version — prompt size plus content size, times price per input token, plus output cost, all times the number of posts scored — gets reduced by crossing out the negligible parts. What’s left is prompt size × price per input token × number of posts. That worked out to about $250 a month with Mixtral 8x7b.

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:

A table comparing a custom ML model with an LLM. The custom model needs a training set, adapts automatically as new labels come in, and is much cheaper — but shipping a new version is hard and it offers no explainability. The LLM needs only an evaluation set, is easier to iterate on and somewhat explainable, but is more expensive and frozen in time, so the prompt and evaluation set need maintaining.

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?

Two notes on our other models. For the topic model, the dataset is free, and switching to an LLM would mean no longer adapting to new trends. For the spam model, an LLM would be expensive because there are so many posts to score, and a static LLM plus a fixed prompt is a problem in a cat-and-mouse game.

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

A Medium story with its topic tags highlighted at the top — writers add topics to their stories and we leverage that to build a training dataset “for free”. The model automatically adapts to new trends as new stories get posted on Medium

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

A Medium story titled “How to Keep a Glass Wine Cellar From Looking Like a Display Case,” tagged Architecture, Wine and Design. The writing is bland and generic, and the first sentence contains a link — some freshly caught spam (SEO hacking)

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.

The same six stories ranked by NSFW score, shown with two different thresholds. At a threshold of 90, only the top story is flagged, so two genuinely NSFW stories slip through as false negatives. At a threshold of 30, almost everything is flagged, so a perfectly SFW story gets caught as a false positive and won’t be distributed. Scores let you move that line wherever you want.

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

A results table with three columns: story title, prediction, and reason. Four stories scored 90 or 95, each with a one-line explanation from the model — “explicit sexual content in one”, “selling illegal drugs online”, “erotic content” for the two other examples. — spot checking and debugging is easier with a short explanation

A bar chart, “Distribution of NSFW posts by high level category.” “Explicit sexual content” is by far the largest share at around 37%, followed by “violence and gore” at roughly 12% and “erotica” at about 10%. “Drugs”, “pornography” and “hate speech” each account for a few percent or less. — asking for a description lets you break things down nicely

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.