Model component API improvements

[favyen2]

Model component API improvements.

• Allow intermediate outputs (passed between model components) to be arbitrary, but provide standard types that can be used in most scenarios. This makes it more clear what the expected input and output for each model component is.
• Make the encoder components (after the initial feature extractor) and decoder components (before the final predictor) share the same API.
• Create a new SampleMetadata class that replaces the metadata dict.
• Create a new ModelContext class that includes the input dicts and SampleMetadatas to pass to models instead of just the inputs. This way models can use location/timestamp info if it's useful.
• Add abstract base classes for FeatureExtractor, IntermediateComponent, and Predictor so the signature of the forward pass is clear for all these model components.
• Improve the documentation and type checking for task-specific outputs (output from Predictor that gets split up per-example and eventually passed to Task.process_output).

Potential config-breaking changes:

• I removed the collapse option from PickFeatures. This was previously used to turn a list of BCHW feature maps into one BCHW tensor, since many task heads expect one BCHW tensor as input. But now those heads expect a FeatureMaps instead so collapse is not needed.

[yawenzzzz]

LGTM!

[favyen2]

I tested this in olmoearth_projects with AWF model and the prediction through olmoearth_run worked, just needed to get rid of the PickFeatures which is no longer needed (and no longer supports the collapse argument as mentioned in the PR). It's possible there are bugs with other models but we can fix them as the arise.