Token-cost classifier (ChatGPT, 4 buckets, distilbert-base-uncased)

Fine-tuned from distilbert-base-uncased to predict how many output tokens ChatGPT will generate for a prompt, before that model is run.

Trained on single-turn English WildChat conversations. Labels are output-token counts from the GPT tokenizer (cl100k_base).

The head is 4-way classification over length buckets. Cut points are in bucket_edges.json (pandas cut bins: counts up to each edge inclusive):

Training context length: 256 tokens.

from transformers import AutoModelForSequenceClassification, AutoTokenizer

repo = "Nadi-Boss/token-cost-wildchat-distilbert-4bucket"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSequenceClassification.from_pretrained(repo)
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256)
bucket = int(model(**inputs).logits.argmax(-1))