Refactor latents and prompt creation

[luciaquirke]

The latents code and the prompt creation explainer code could be significantly simplified.

Simplifying the code would facilitate prototyping and adding new features.

E.g.: reshaping the input tokens to be of sequence length window_size before caching the latents would allow us to remove most or all of the window handling logic in constructors.

[SrGonao]

Is there a reason for the simplification you suggested? It seems to solidify stuff more; by e.g. not being able to choose different windows sizes without having to re-cache activations.
Probably agree on the prompt creation code.
What specific simplifications do you envision?

[luciaquirke]

I'm mostly interested in consolidating logic so it's easier to reason about and add to.

The window logic suggestion was because the current logic is kind of hard to follow such that it took me quite a while to pipe through a super simple feature (adding the top k model predictions at a token to an explainer), but if being able to quickly iterate between window sizes is worth the complexity we could keep it, removing would be a solidification that I thought was fine because I didn't know anyone was experimenting with window size.

Other things:

• make it easier to find the entry point for the pipeline in _aprocess_latent, e.g. with comment or maybe a clearer function flow if there is one
• make it easier to understand the calls to constructor and sampler, they currently have ambiguous names and mutate the record in non-obvious ways. I have an intuition that we could merge them into a single create_record function and then simplify the logic because we won't need to form an intermediate version of the record object
• the constructor + sampler function could potentially then become its own pipeline component after the loader and before the explainers, making it possible to interact with the record creation logic in isolation from the loading logic and vice versa