Refactor press implementation

[maxjeblick]

Feature

Separate press class into two separate classes

• A scorer class that implements the .score method
• A pruning class that implements the .forward_hook method

The press class then works with dependency injection, e.g., ExpectedAttentionPress can be expressed as

press = BasePruner(
        compression_ratio=compression_ratio,
        scorer=ExpectedAttentionScorer(
            n_future_positions=n_future_positions, n_sink=n_sink, use_covariance=use_covariance, use_vnorm=use_vnorm
        ),
    )

Motivation

Current press code couples forward hook and score method, making it harder to implement custom workflows.
By decoupling pruning and scoring functionality, it is possible to add new pruning methods by subclassing BasePruner rather than using a wrapper function (e.g. PerLayerCompressionPruner).

[SimJeg]

Other features I can think about for a v0.1.0:

• a utils.py module centralizing functions that are repeatedly use across different presses (e.g. compute_queries)
• a distinction between SequenceScorer and ChannelScorer if more presses like ThinkPress appear (see Add ThinKPress #20)
• a refactored of the README to move away from centralizing everything around the compression_ratio parameter. For some methods, the compression depends on the prompts, hence the compression ratio is not an input to the press but an output

[maxjeblick]

a distinction between SequenceScorer and ChannelScorer

Could be implemented as follows:

sequence_press = ... (normal press)
press = ChannelPruner(compression_ratio=..., sequence_press=sequence_press, channel_scorer=ThinkChannelScorer(....))

with ThinkChannelScorer having a score_channel method (instead of score).