♻️ refactor wandb import

"""

Callback functions that can be invoked while fitting a KerasModel.

"""

import tensorflow as tf

import sys

from deepchem.models.keras_model import is_wandb_available

if is_wandb_available():

  import wandb

class ValidationCallback(object):

  """Performs validation while training a KerasModel.

      for key in scores:

        model._log_value_to_tensorboard(tag=key, simple_value=scores[key])

    if model.wandb:

      import wandb

      wandb.log(scores, step=step)

    if self.save_dir is not None:

      score = scores[self.metrics[self.save_metric].name]

except:

  from collections import Sequence as SequenceCollection

logger = logging.getLogger(__name__)

from deepchem.data import Dataset, NumpyDataset

from deepchem.metrics import Metric

from deepchem.models.losses import Loss

except (ImportError, AttributeError):

  _has_wandb = False

def is_wandb_available():

  return _has_wandb

logger = logging.getLogger(__name__)

class KerasModel(Model):

    self.tensorboard = tensorboard

    # W&B logging

    if wandb and not is_wandb_available():

    if wandb and not _has_wandb:

      logger.warning(

          "You set wandb to True but W&B is not installed. To use wandb logging, "

          "run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface."

      )

    self.wandb = wandb and is_wandb_available()

    self.wandb = wandb and _has_wandb

    # Backwards compatibility

    if "tensorboard_log_frequency" in kwargs:

    avg_loss = 0.0

    last_avg_loss = 0.0

    averaged_batches = 0

    train_op = None

    if loss is None:

      loss = self._loss_fn

    var_key = None

    variances: Optional[List[np.ndarray]] = None

    if (outputs is not None) and (other_output_types is not None):

      raise ValueError(

          'This model cannot compute outputs and other output_types simultaneously. Please invoke one at a time.'

      )

          'This model cannot compute outputs and other output_types simultaneously.'

          'Please invoke one at a time.')

    if uncertainty and (other_output_types is not None):

      raise ValueError(

          'This model cannot compute uncertainties and other output types simultaneously. Please invoke one at a time.'

      )

          'This model cannot compute uncertainties and other output types simultaneously.'

          'Please invoke one at a time.')

    if uncertainty:

      assert outputs is None

      if self._variance_outputs is None or len(self._variance_outputs) == 0:

    if (outputs is not None and self.model.inputs is not None and

        len(self.model.inputs) == 0):

      raise ValueError(

          "Cannot use 'outputs' argument with a model that does not specify its inputs. Note models defined in imperative subclassing style cannot specify outputs"

          "Cannot use 'outputs' argument with a model that does not specify its inputs."

          "Note models defined in imperative subclassing style cannot specify outputs"

      )

    if isinstance(outputs, tf.Tensor):

      outputs = [outputs]

import numpy as np

import torch

import torch.utils.tensorboard

import time

import logging

import os

except:

  from collections import Sequence as SequenceCollection

logger = logging.getLogger(__name__)

from deepchem.data import Dataset, NumpyDataset

from deepchem.metrics import Metric

from deepchem.models.losses import Loss

except (ImportError, AttributeError):

  _has_wandb = False

def is_wandb_available():

  return _has_wandb

logger = logging.getLogger(__name__)

class TorchModel(Model):

    self.model.to(device)

    # W&B logging

    if wandb and not is_wandb_available():

    if wandb and not _has_wandb:

      logger.warning(

          "You set wandb to True but W&B is not installed. To use wandb logging, "

          "run `pip install wandb; wandb login` see https://docs.wandb.com/huggingface."

      )

    self.wandb = wandb and is_wandb_available()

    self.wandb = wandb and _has_wandb

    self.log_frequency = log_frequency

    if self.tensorboard: