[Question]: How to obtain the masked and kept tokens?

[jamesnulliu]

Describe the issue

I noticed that in PromptCompressor.get_compressed_input, there is a variable called need_idx. Is that all I need?

[natalienma]

From my understanding, you can prune tokens in multiple ways, but in the paper, they create a dataset of pairs of original text and compressed text by prompting GPT-4 to remove unnecessary tokens. They then train a small BERT-style encoder on this dataset. During inference, BERT scores each token in the prompt either a 0 (drop), or 1 (keep).

In my project, I introduced the idea of soft scoring, which makes the scores continuous between 0 and 1, and you can toggle how much content you want to retain -- anything scoring above a 0.7, for example.

I tried to soft score by prompting GPT-4 to score each word from 0.0-1.0, providing more nuance for pruning. The failure mode I experienced: GPT-4 hallucinates due to lack of reasoning. This hallucination compounded over time and soon all scores were converging to 0.

Happy to share more details on the setup if useful.