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 has two outputs, (mu_z, log_sigma_z), for a heteroscedastic Gaussian on log1p(output_tokens). Those outputs are in a standardized space. target_stats.json holds the training mean and standard deviation of log1p(count) used to convert back to token space. The median token count is expm1(mu_log).
Training context length: 256 tokens.
import math
import json
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForSequenceClassification, AutoTokenizer
repo = "Nadi-Boss/token-cost-llama8b-distilbert-regression"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSequenceClassification.from_pretrained(repo)
stats = json.load(open(hf_hub_download(repo, "target_stats.json")))
inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=256)
with torch.no_grad():
mu_z, log_sigma_z = model(**inputs).logits[0].tolist()
mu_log = mu_z * stats["std"] + stats["mean"]
sigma_log = math.exp(log_sigma_z) * stats["std"]
median_tokens = math.expm1(mu_log)