text: """WILDS integration for Meerkat."""
  from __future__ import annotations
  
  from argparse import Namespace
  from typing import List
  
  import numpy as np
  import pandas as pd
  from datasets import DatasetInfo
  from torch.utils.data._utils.collate import default_collate
  
  from meerkat.columns.abstract import AbstractColumn
  from meerkat.columns.numpy_column import NumpyArrayColumn
  from meerkat.datapanel import DataPanel
  from meerkat.tools.identifier import Identifier
  
  from .config import base_config, populate_defaults
  from .transforms import initialize_transform
  
  try:
      import wilds
      from wilds.datasets.wilds_dataset import WILDSSubset
  
      _wilds_available = True
  except ImportError:
      _wilds_available = False
  
  
  def get_wilds_datapanel(
      dataset_name: str,
      root_dir: str,
      version: str = None,
      identifier: Identifier = None,
      column_names: List[str] = None,
      info: DatasetInfo = None,
      split: str = None,
      use_transform: bool = True,
      include_raw_input: bool = True,
  ):
      """Get a DataPanel that holds a WildsInputColumn alongside NumpyColumns for
      targets and metadata.
  
      Example:
      Run inference on the dataset and store predictions alongside the data.
      .. code-block:: python
  
          dp = get_wilds_datapane("fmow", root_dir="/datasets/", split="test")
          model = ... # get the model
          model.to(0).eval()
  
          @torch.no_grad()
          def predict(batch: dict):
  
              out = torch.softmax(model(batch["input"].to(0)), axis=-1)
              return {"pred": out.cpu().numpy().argmax(axis=-1)}
  
          dp = dp.update(function=predict, batch_size=128, is_batched_fn=True)
  
  
      Args:
  
          dataset_name (str, optional): dataset name. Defaults to `"fmow"`.
          version (str, optional): dataset version number, e.g., '1.0'.
              Defaults to the latest version.
          root_dir (str): the directory where the WILDS dataset is downloaded.
              See https://wilds.stanford.edu/ for download instructions.
          split (str, optional): see . Defaults to None.
          use_transform (bool, optional): Whether to apply the transform from the
              WILDS example directory on load. Defaults to True.
          identifier (Identifier, optional): [description]. Defaults to None.
          column_names (List[str], optional): [description]. Defaults to None.
          info (DatasetInfo, optional): [description]. Defaults to None.
          use_transform (bool, optional): [description]. Defaults to True.
          include_raw_input (bool, optional): include a column for the input without
              the transform applied – useful for visualizing images. Defaults to True.
      """
      if not _wilds_available:
          raise ImportError(
              "The WILDS package is not installed. To use the RG WILDS module, please "
              "install by following instructions at "
              "https://wilds.stanford.edu/get_started/"
          )
  
      input_column = WILDSInputColumn(
          dataset_name=dataset_name,
          version=version,
          root_dir=root_dir,
          split=split,
          use_transform=use_transform,
      )
      output_column = input_column.get_y_column()
      metadata_columns = input_column.get_metadata_columns()
  
      data = {"input": input_column, "y": output_column, **metadata_columns}
      if include_raw_input:
          data["raw_input"] = input_column.copy()
          data["raw_input"].use_transform = False
          data["raw_input"]._data.transform = lambda x: x
          data["raw_input"]._collate_fn = lambda x: x
  
      return DataPanel(
          data,
          identifier=identifier,
          column_names=column_names,
          info=info,
          split=split,
      )
  
  
  class WILDSInputColumn(AbstractColumn):
      def __init__(
          self,
          dataset_name: str = "fmow",
          version: str = None,
          root_dir: str = None,
          split: str = None,
          use_transform: bool = True,
          **kwargs,
      ):
          """A column wrapper around a WILDS dataset that can lazily load the
          inputs for each dataset.
  
          Args:
              dataset_name (str, optional): dataset name. Defaults to `"fmow"`.
              version (str, optional): dataset version number, e.g., '1.0'.
                  Defaults to the latest version.
              root_dir (str, optional): the directory  . Defaults to None.
              split (str, optional): the split . Defaults to None.
              use_transform (bool, optional): Whether to apply the transform from the
                  WILDS example directory on load. Defaults to True.
          """
  
          self.dataset_name = dataset_name
          self.version = version
          self.root_dir = root_dir
          self.split = split
          self.use_transform = use_transform
  
          dataset = wilds.get_dataset(
              dataset=dataset_name, version=version, root_dir=root_dir
          )
          self.root = dataset.root
          self.split = split
  
          self.metadata_columns = {}
          # get additional, dataset-specific metadata columns
          if dataset_name == "fmow":
              metadata_df = dataset.metadata
              metadata_df = metadata_df[metadata_df.split != "seq"]
              if self.split is not None:
                  metadata_df = metadata_df[
                      dataset.split_array == dataset.split_dict[self.split]
                  ]
              self.metadata_columns.update(
                  {
                      field: NumpyArrayColumn(data=series.values)
                      for field, series in metadata_df.iteritems()
                  }
              )
  
          if use_transform:
              # we need a WILDS config in order to initialize transform
              config = Namespace(dataset=dataset_name, **base_config)
              config = populate_defaults(config)
              transform_name = (
                  config.train_transform if split == "train" else config.eval_transform
              )
              transform = initialize_transform(
                  transform_name, config=config, dataset=dataset
              )
              # wilds defaults to torch `default_collate`
              collate = default_collate if dataset.collate is None else dataset.collate
          else:
              transform = collate = None
  
          if split is not None:
              dataset = dataset.get_subset(split, transform=transform, frac=1.0)
          elif transform is not None:
              # only WILDSSubset supports applying a transform, so we use it even if no
              # split is applied
              dataset = WILDSSubset(dataset, np.arange(len(dataset)), transform=transform)
  
          self.metadata_columns.update(
              {
                  f"meta_{field}": NumpyArrayColumn(data=dataset.metadata_array[:, idx])
                  for idx, field in enumerate(dataset.metadata_fields)
              }
          )
          super(WILDSInputColumn, self).__init__(data=dataset, collate_fn=collate)
  
      def get_y_column(self):
          """Get a NumpyArrayColumn holding the targets for the dataset.
  
          Warning: `WildsDataset`s may remap indexes in arbitrary ways so it's important
          not to directly try to access the underlying data structures, instead relying on
          the `y_array` and `metadata_array` properties which are universal across WILDS
          datasets.
          """
          return NumpyArrayColumn(data=self.data.y_array)
  
      def get_metadata_columns(self):
          return self.metadata_columns
  
      def _get_cell(self, index: int, materialize: bool = True):
          # only get input (not y and meta)
          return self.data[index][0]
  
      def _repr_pandas_(
          self,
      ) -> pd.Series:
          series = pd.Series(np.arange(len(self._data)))
          return series.apply(
              lambda x: f"WildsInput(path={self.root}/images/rgb_img_{x}.png)"
          )
  
      @classmethod
      def _state_keys(cls) -> set:
          """List of attributes that describe the state of the object.
  
          Warning: this write
          is very lightweight, only the name of the dataset (`dataset_name`), the
          directory of the dataset `root_dir`, and the other args to `__init__` are
          written to disk. If the data at `root_dir` is modified, `read` will return a
          column with different data.
          """
          return {
              "dataset_name",
              "version",
              "root_dir",
              "split",
              "use_transform",
              "_visible_rows",
          }
  