Update per token reward

[ljvmiranda921]

Closes #22

Kinda big PR that adds two new scripts:

• get_per_token_reward that computes the per-substring reward and then keeps it in a directory. The name of the directory is the hashed value of the input text.
• draw_per_token_reward accepts a token hash and draws the visualization as seen below. I just made them into a single heatmap so that the spacing is easier to manage. Note that the colorbar spacing might look different as we add new models.

How this works

python -m analysis.per_token_reward "hi there general kenobi" --model IDEA-CCNL/Ziya-LLaMA-7B-Reward

Say the hashed value of the text "hi there general kenobi" is 71b1349c9b, it will then create a directory abcdefgh123 with a filename based on the combined hash of the model name and chat_template:

71b1349c9b/
    59761c2c8a.json

And then the JSON file looks like this:

{
    "text": "hi there general kenobi",
    "text_hash": "71b1349c9b",
    "model": "IDEA-CCNL/Ziya-LLaMA-7B-Reward",
    "chat_template": "tulu",
    "model_chat_hash": "59761c2c8a",
    "substrings": [
        "hi",
        "hi there",
        "hi there general",
        "hi there general k",
        "hi there general ken",
        "hi there general kenobi"
    ],
    "tokens": [
        "hi",
        "there",
        "general",
        "k",
        "en",
        "obi"
    ],
    "rewards": [
        -1.9443359375,
        -1.1064453125,
        0.26318359375,
        1.2578125,
        1.208984375,
        3.24609375
    ]
}

Ideally we store this in a file somewhere where we have a directory of text prompts (and model rewards in JSON). To create the visualization, we use the script draw_per_token_reward. We pass the hash and a path to the output PNG file:

python -m analysis.draw_per_token_reward 71b1349c9b test.png

[ljvmiranda921]

Not sure if this is the best approach. I manually listed possible reward models per "group." I feel like the initial approach that looks for substrings might introduce some hidden errors. In this case, it will only error out if it's not added to models.

[natolambert]

Yeah, tokenizers gonna tokenizer. If you're happy with it, it's okay for now.

[natolambert]

Some general notes as I go through:

1. I also think a line plot could be cool. We can iterate on that if anyone is interested in the topic further :)
2. What happens with Per token multiple rms #29 , did y'all coordinate on this?
3. Regarding tokenizer, if it's a line plot we could just plot per character (and some tokens jump multiple on the x axis), which may make some of the code less painful
4. The model config stuff can maybe be a separate PR. We can just expand it after this is merged.
5. I'm a little wary of spacy-alignments package. Not well used, and I'm not convinced we need to go in that direction. Lmk what you think.

Anyways, approving if you feel strongly about it, but we'll prolly discuss more before merging.

[ljvmiranda921]

Thanks for the review!

1. Sure, I can try a line chart, but might be able to start working on it the latter half of this week 👍
2. Ah yeah I used it as reference for this PR. I think we can close Per token multiple rms #29 once this PR is merged, any thoughts @khyathiraghavi ?
3. Sure! Same as Multiple styles of computing reward with DPO #1
4. Hmm, right now it's tightly coupled so it might be hard to separate them. We can merge this first, then I can make another PR that separates them (or maybe that's what you meant in the first place?).
5. I initially wrote a naive if substring in token logic but I encountered many edge-cases. I looked around and found that this is the most robust so far. If you're still wary, we can put the dependency in requirements.txt since spacy-alignments is not used in the core training scripts.

[khyathiraghavi]

Yes, we coordinated and @ljvmiranda921 merged my PR https://github.com/allenai/herm/tree/per-token-multiple-rms in the current version, it looks good to me!
We can also consider a good RM as a base and use that tokenizer in place of white space and compare the rest against this in the future -- the current version looks good to me for now.