1from transformers import pipeline
2
3classifier = pipeline(
4 "text-classification",
5 model="mrdbourke/ettin-150m-food-or-drink-classifier",
6 device="cuda", # or "cpu"
7 torch_dtype="float16" # for faster inference
8)
9
10# Single prediction
11result = classifier("A bowl of ramen with soft-boiled egg and nori")
12print(result)
13# [{'label': 'food_or_drink', 'score': 0.9995}]
14
15# Batch prediction
16texts = [
17 "A glass of red wine next to a cheese board",
18 "A yellow tractor driving over a grassy hill",
19 "Fresh squeezed orange juice with ice",
20 "A laptop computer on a wooden desk",
21]
22results = classifier(texts, batch_size=512)
23for text, r in zip(texts, results):
24 print(f"{r['label']:<20s} {r['score']:.4f} {text}")
1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3
4tokenizer = AutoTokenizer.from_pretrained("mrdbourke/ettin-150m-food-or-drink-classifier")
5model = AutoModelForSequenceClassification.from_pretrained(
6 "mrdbourke/ettin-150m-food-or-drink-classifier",
7 torch_dtype=torch.float16,
8).to("cuda").eval()
9
10texts = ["A bowl of ramen", "A red car on the highway"]
11inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to("cuda")
12
13with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
14 outputs = model(**inputs)
15 probs = torch.softmax(outputs.logits.float(), dim=-1)
16 preds = torch.argmax(probs, dim=-1)
17
18labels = ["food_or_drink", "not_food_or_drink"]
19for text, pred, prob in zip(texts, preds, probs):
20 print(f"{labels[pred]:<20s} {prob[pred]:.4f} {text}")
The teacher model (ModernBERT-large zero-shot NLI) processes ~871 rows/s because it requires encoding hypothesis pairs for each label. The fine-tuned student model does a single forward pass with a classification head, achieving 5,874+ rows/s — a 7x speedup that makes billion-scale inference practical.
