How to restrict the model to only consider a set of possible tokens

Hi, how do I make openai.Completion.create() calls to only consider a small set of tokens as possible so that the choices and most likely tokens will always be among these small set of tokens?

I’m using the completion create feature to study computational models of sequence learning and I need to restrict to only 8 possible tokens since the probabilistic sequences I want to give to the model as a prompt have only 8 tokens.

You could try using logit_bias.

Just be sure to select the correct tokenizer (and tokens) for the model when defining your map.

1 Like

Hi @curt.kennedy thanks so much for offering your help here. Unfortunately, this is not working for me.

Here is my code:

MODEL = "davinci"
enc = tiktoken.encoding_for_model(MODEL)
logit_bias_map = {}
prompt = " john tom dan john bob dan john tom will john tom"
for x in prompt.split():
    logit_bias_map[str(enc.encode(x)[0])] = 100
completions = openai.Completion.create(
    logit_bias = logit_bias_map

And here is the ouput:

[<OpenAIObject at  JSON: {
  "finish_reason": "length",
  "index": 0,
  "logprobs": {
    "text_offset": [
    "token_logprobs": [
    "tokens": [
      " tom",
      " dan",
      " john",
      " bob",
      " dan",
      " john",
      " tom",
      " will",
      " john",
      " tom",
    "top_logprobs": [
        ",": -3.6765778,
        "-": -2.8618271,
        ".": -3.5472288,
        "bytes:\\xe2\\x80": -3.9375181,
        "s": -3.2035592
        "as": -3.009693,
        "ase": -2.301611,
        "asi": -2.3654172,
        "lin": -1.7576712,
        "my": -2.7895849
        "\n": -4.324432,
        " and": -3.30854,
        " j": -2.8703802,
        "forth": -3.821676,
        "ley": -3.9747922
        "\n": -3.7143948,
        " dan": -3.0179114,
        " tom": -2.18535,
        "ny": -3.0394833,
        "son": -2.2874274
        "\n": -3.0000987,
        " bob": -3.4359412,
        " dan": -3.138067,
        " john": -2.964168,
        ")": -3.2884574
        "\n": -3.6697154,
        " bob": -1.0042298,
        " dan": -3.6627877,
        " john": -2.0533974,
        " tom": -3.092765
        "\n": -3.4838939,
        " bob": -1.1732672,
        " dan": -3.6677318,
        " john": -3.3311195,
        " tom": -1.9761846
        "\n": -3.5157578,
        " bob": -1.8858293,
        " dan": -0.95853364,
        " john": -3.0290709,
        " tom": -2.9475436
        " be": -2.5264819,
        " dan": -3.6445656,
        " john": -1.993683,
        "iam": -1.3272516,
        "y": -4.225769
        " bob": -1.0446652,
        " dan": -2.8825402,
        " john": -3.5926692,
        " tom": -1.1941555,
        " will": -2.8907952
        "b": -2.1100414,
        "dan": -1.2277858,
        "john": -1.8985239,
        "tom": -2.3150966,
        "will": -1.0868405
  "text": "john tom dan john bob dan john tom will john tomwill"

And if we print that logit_bias_map we get one token per each name in the prompt as shown by the 5 tokens for the 5 names:

{'30686': 100, '39532': 100, '25604': 100, '65': 100, '10594': 100}

What doesn’t seem right to me is that the top_logprobs has different tokens other than the 5 I set to 100 (and they are not super low either in their logprobs). What do you think? Seems like something is off, and wondering if there’s a solution path you can help point me to.

Thanks for your help!

hi @curt.kennedy , I would appreciate yours or other experts thoughts on the above. Perhaps if it is not you, then there is someone part of the team who is an expert on the logit_bias_map that can offer some advice? I’m wondering if there are particular models that I can use which have the setting of a token as 100 be an exclusive selection while others do not allow for exclusive selection using logit_bias_map set to 100 (the models I tried belonging to the latter group, explaining my inability to achieve exclusive selection)?

Hi, still scratching my head on this one. Been doing a bunch of business travel lately, so haven’t had time to look deeper, and provide an example.

My one question though, is that it seems like your prompt of just listing the names is not directing DaVinci to do anything. So wondering if that is part of the problem? Maybe a more realistic example would help me here.

Also, if you try lowering your value of 100 to 30 or something lower. The higher values might still be glitched.

Also, try removing your leading space.

I think it’s an interesting question to see if token biasing will “restrict” the model to those tokens. The thing is, you are restricting an unrestricted (non fine-tuned) model, so its search space is huge. You likely need some steering towards your desired tokens, to increase the odds of it landing on your desired tokens.

Fine-tuning should shrink the search space, getting closer to your tokens, but you need to train it. This is how I normally see the “small output token” option, is with a fine-tune, not a biased token unconstrained model.

P.S. Another tip to coerce the model into your tokens, is at least constraining the token count in the output. I see you already did this, but to the layperson, you can get the model to self-constrain and more likely align with what you are thinking. For example, I sent the unconstrained and constrained version of this to DaVinci:

“Because I like eating cheese, I also like eating”

With the unconstrained, it went on and on about eating Mac and Cheese. This was not my desired and expected short answer. I was expecting " crackers" as the completion and that is what I got when I restrict to a 2 token output, as shown here:

It looks like you were already aware of this, but emphasizing for the general audience.

Another option to restrict tokens is prompt engineering with the higher end models (works on GPT-4, but not Turbo for me). I didn’t bias anything and am using the foundation model (not a fine-tune). :sunglasses:

(Edit: Maybe nix the above prompt, because it is just repeating Bob, Dan, Curt in a loop. I think lesson learned is that if it doesn’t make sense to a human, it certainly won’t make sense to AI).

OK, last edit, promise :crazy_face: Here I am guiding the AI to specific output tokens and asking it to make sense of things. It’s doing a good job now, but you have to “give it the framework” to do so in the prompt.

Is the crackers images from the playground?

Yes with “Full Spectrum” selected in Playground for DaVinci.

1 Like