Token-cost regression (WildChat, DistilBERT)

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