trlx: (perhaps) data preparation bug in RLHF

The hack has produced a similar error:

KeyError: "SUBREDDIT: r/relationships\nTITLE: [Update 2] I [18 M] want to ask out a girl [18 F] out on a date, general tips needed.\nPOST: [Original](\n(Clarification on this one, I didn’t mean the one as the girl I wanted to marry)\nTL;DR: "

I’ve reproduced this problem before. The temporary (terrible) hack I used was to truncate the OpenAI summary keys to 272 chars. Something like:

    # NOTE: `num_char_compares` is a hacky way to avoid `KeyError`s when retrieving original summaries because we can't
    # ensure the decoded summaries are the same as the original summaries (non-invertibility of the used tokenizers?)
    num_char_compares = 272  # Only compare the first N chars in the decoded post for safe mapping

    def reward_fn(samples: List[str], **kwargs):
        original_samples = [text.split("TL;DR:")[0] + "TL;DR: " for text in samples]
        # Truncate text before accessing `post_summary_dict`
        original_samples = [text + post_summary_dict[text[:num_char_compares]] for text in original_samples]
        original_scores = get_scores(original_samples)
        scores = get_scores(samples)
        norms_scores = scores - original_scores
        return norms_scores

    # ...

    # Get the OpenAI summaries
    post_summary_dict = {}
    train_prompts = get_prompt_dataset(train_posts, max_length_input)
    for i in range(len(train_prompts)):
        post_summary_dict[train_prompts[i][:num_char_compares]] = train_summaries[i]
    val_prompts = get_prompt_dataset(val_posts, max_length_input)
    for i in range(len(val_prompts)):
        post_summary_dict[val_prompts[i][:num_char_compares]] = val_summaries[i]

Still looking for a proper fix…