Enhance WILDSSubset to support flexible data returns

[rdarbinyan]

Summary

This update modifies the WILDSSubset class to support flexible return types from its __getitem__ method, accommodating both standard and custom dataset requirements. The enhancement allows for varied data returns, from basic x and y pairs to more complex structures including additional information like background labels. Importantly, this change is backward compatible, ensuring existing datasets continue to function as before.

Changes

• Enhanced __getitem__ to dynamically handle varying numbers of returned items, enabling support for datasets that return additional data elements or those that simplify to just x and y.
• Maintained backward compatibility with standard x, y, metadata format datasets, while also facilitating custom datasets with expanded or reduced data structures.

Motivation

The update is motivated by the need for greater flexibility in dataset handling within the WILDS library. A personal use case involved inheriting from WaterbirdsDataset and making it return x, y, s, and metadata. In another scenario, a user might prefer a dataset to return only the features (x) and labels (y), omitting the metadata for a leaner data structure. This change facilitates such customizations by allowing easy adjustments to the return types from dataset classes, without needing to alter the classes directly.

Conclusion

By introducing this backward-compatible enhancement, the WILDS library becomes more adaptable and user-friendly, encouraging contributions and customizations that meet diverse research and application needs.