text: # Copyright The PyTorch Lightning team.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  from typing import Any, Callable, Dict, List, Optional, Union
  
  import torchmetrics
  
  from flash.core.adapter import Adapter
  from flash.core.data.io.input import DataKeys
  from flash.core.model import Task
  
  
  class PytorchTabularAdapter(Adapter):
      def __init__(self, backbone):
          super().__init__()
  
          self.backbone = backbone
  
      @classmethod
      def from_task(
          cls,
          task: Task,
          task_type,
          embedding_sizes: list,
          categorical_fields: list,
          cat_dims: list,
          num_features: int,
          output_dim: int,
          backbone: str,
          loss_fn: Optional[Callable],
          metrics: Optional[Union[torchmetrics.Metric, List[torchmetrics.Metric]]],
          backbone_kwargs: Optional[Dict[str, Any]] = None,
      ) -> Adapter:
  
          backbone_kwargs = backbone_kwargs or {}
          parameters = {
              "embedding_dims": embedding_sizes,
              "categorical_cols": categorical_fields,
              "categorical_cardinality": cat_dims,
              "categorical_dim": len(categorical_fields),
              "continuous_dim": num_features - len(categorical_fields),
              "output_dim": output_dim,
          }
          adapter = cls(
              task.backbones.get(backbone)(
                  task_type=task_type, parameters=parameters, loss_fn=loss_fn, metrics=metrics, **backbone_kwargs
              )
          )
  
          return adapter
  
      @staticmethod
      def convert_batch(batch):
          new_batch = {
              "continuous": batch[DataKeys.INPUT][1],
              "categorical": batch[DataKeys.INPUT][0],
          }
          if DataKeys.TARGET in batch:
              new_batch["target"] = batch[DataKeys.TARGET].reshape(-1, 1)
          return new_batch
  
      def training_step(self, batch, batch_idx) -> Any:
          return self.backbone.training_step(self.convert_batch(batch), batch_idx)
  
      def validation_step(self, batch, batch_idx):
          return self.backbone.validation_step(self.convert_batch(batch), batch_idx)
  
      def test_step(self, batch, batch_idx):
          return self.backbone.validation_step(self.convert_batch(batch), batch_idx)
  
      def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
          return self(self.convert_batch(batch))
  
      def forward(self, batch: Any) -> Any:
          return self.backbone(batch)["logits"]
  