Repository: stanford-crfm/helm
Issue: Refactor WindowService
GitHub labels: p2

framework

cleanup

The problems with the current system are that every time we want to add a model with a new sequence length for an existing model provider, the user has to:

- Create a new [CONTEXT_WINDOW_TAG](https://github.com/stanford-crfm/helm/blob/d69a6e696e156ec175a4b6d33e18faf4e2686a7e/src/helm/proxy/models.py#L22)
- Create a new `WindowService` that sets the new sequence length
- Modify framework code in [`WindowServiceFactory.get_window_service`](https://github.com/stanford-crfm/helm/blob/d69a6e696e156ec175a4b6d33e18faf4e2686a7e/src/helm/benchmark/window_services/window_service_factory.py#L45)

The issues here are:

- Proliferation of model tags.
- It's hard to tell the sequence length of each model, and easy to introduce bugs by introducing the wrong sequence length.
- The huge if-else branch tree in `WindowServiceFactory` is not user-friendly, and is a potential source of bugs.

The proposed refactor:

Currently we have [`WindowService`](https://github.com/stanford-crfm/helm/blob/d69a6e696e156ec175a4b6d33e18faf4e2686a7e/src/helm/benchmark/window_services/window_service.py#L24) has a ton of stuff:

- Model metadata: `tokenizer_name`, `max_sequence_length`, `max_request_length`
- Tokenizer metadata: `end_of_text_token`, `prefix_token`
- Tokenizer operations: `tokenize`, `encode`, `decode`, `get_num_tokens`
- Windowing logic: `fits_within_context_window`, `truncate_from_right`

We should break these out of `WindowService` as follows:

- Move model metadata to the model: `tokenizer_name`, `max_sequence_length`, `max_request_length`
- Move tokenizer metadata to the tokenizer: `end_of_text_token`, `prefix_token`
- Tokenizer operations are tricky because some of these do additional stuff before passing to the underlying tokenizer.

So basically `WindowService` should only contain the windowing logic. This should also allow us to unify most of the `WindowService`s that are just thin wrappers around the same logic.


This is big refactor so we should do it in steps. #1455 is just for handling `max_sequence_length` and `max_request_length`.

The plan is to introduce a new `Tokenizer` object that would look like this:
```python
class Tokenizer(ABC):
    @property
    @abstractmethod
    def name(self) -> str:
        """The name of the tokenizer."""
        pass

    @property
    @abstractmethod
    def end_of_text_token(self) -> str:
        """The end of text token."""
        pass

    @property
    @abstractmethod
    def prefix_token(self) -> str:
        """The prefix token"""
        pass

    @abstractmethod
    def encode(self, text: str, truncation: bool = False, max_length: Optional[int] = None) -> EncodeResult:
        """Encodes the input text to tokens given the model"""
        pass

    @abstractmethod
    def decode(self, tokens: List[TokenizationToken], normalized_text: Optional[str] = None) -> str:
        """
        Given the model and a list of tokens, outputs the corresponding text.

        For models using the GPT-2 tokenizer, the tokens are integers; for AI21
        models, the tokens are `TokenizationToken`s.

        Some tokenizers (e.g. AI21) normalize the text before encoding it and
        thus require the `normalized_text` for decoding.
        """
        pass

    @abstractmethod
    def tokenize(self, text: str) -> List[str]:
        """Tokenizes the text."""
        pass

    @abstractmethod
    def get_num_tokens(self, text: str) -> int:
        """Tokenizes the text and counts the number of tokens."""
        pass
```

This is _more or less_ done by #1903 and #1943.

There is some remaining cleanup (e.g. remove the `tokenize` method from `WindowService`) but that is relatively minor.