Token-cost classifier (Llama 3.1 8B, 4 buckets, distilbert-base-uncased)

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

Prompts come from LMSYS-Chat-1M. Labels are the output-token counts of Llama 3.1 8B answers.

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-llama8b-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))