Add hooks

Use run-time hooking: This MR introduces a function hooked that can be used like this:

input_ids = ...
hooked_model = hooked(model)
logits = hooked_model(input_ids)

This works in exactly the same way as a non-hooked model, but now you can also access activations within the model (and even change them).