delete pytorch dir for chemberta-tutorial branch

import torch

import torch.nn as nn

from torch.nn import CrossEntropyLoss, MSELoss

from transformers.modeling_roberta import (

    ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,

    BertPreTrainedModel,

    RobertaClassificationHead,

    RobertaConfig,

    RobertaModel,

)

class ChemBERTaforSequenceClassification(BertPreTrainedModel):

    r"""

        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size,)``:

            Labels for computing the sequence classification/regression loss.

            Indices should be in ``[0, ..., config.num_labels]``.

            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),

            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:

        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:

            Classification (or regression if config.num_labels==1) loss.

        **logits**: ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``

            Classification (or regression if config.num_labels==1) scores (before SoftMax).

        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)

            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)

            of shape ``(batch_size, sequence_length, hidden_size)``:

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.

        **attentions**: (`optional`, returned when ``config.output_attentions=True``)

            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

    Examples::

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')

        model = ChemBERTaforSequenceClassification.from_pretrained('roberta-base')

        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1

        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1

        outputs = model(input_ids, labels=labels)

        loss, logits = outputs[:2]

    """  # noqa: ignore flake8"

    config_class = RobertaConfig

    pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST

    base_model_prefix = "roberta"

    def __init__(self, config, weight=None):

        super(ChemBERTaforSequenceClassification, self).__init__(config)

        self.num_labels = config.num_labels

        self.roberta = RobertaModel(config)

        self.classifier = RobertaClassificationHead(config)

        self.weight = weight

    def forward(

        self,

        input_ids=None,

        attention_mask=None,

        token_type_ids=None,

        position_ids=None,

        head_mask=None,

        inputs_embeds=None,

        labels=None,

    ):

        outputs = self.roberta(

            input_ids,

            attention_mask=attention_mask,

            token_type_ids=token_type_ids,

            position_ids=position_ids,

            head_mask=head_mask,

        )

        sequence_output = outputs[0]

        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]

        if labels is not None:

            if self.num_labels == 1:

                #  We are doing regression

                loss_fct = MSELoss()

                loss = loss_fct(logits.view(-1), labels.view(-1))

            else:

                loss_fct = CrossEntropyLoss(weight=self.weight)

                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)

# BELOW code is taken from modles.py methods, for basic idea of structure to follow.

  def fit(self, dataset, nb_epoch=10, batch_size=32, **kwargs):

    """

    Fits a model on data in a Dataset object.

    """

    # TODO(rbharath/enf): We need a structured way to deal with potential GPU

    #                     memory overflows.

    for epoch in range(nb_epoch):

      log("Starting epoch %s" % str(epoch + 1), self.verbose)

      losses = []

      for (X_batch, y_batch, w_batch,

           ids_batch) in dataset.iterbatches(batch_size):

        losses.append(self.fit_on_batch(X_batch, y_batch, w_batch))

      log("Avg loss for epoch %d: %f" % (epoch + 1, np.array(losses).mean()),

          self.verbose)

  def predict(self, dataset, transformers=[], batch_size=None):

    """

    Uses self to make predictions on provided Dataset object.

    Returns:

      y_pred: numpy ndarray of shape (n_samples,)

    """

    y_preds = []

    n_tasks = self.get_num_tasks()

    ind = 0

    for (X_batch, _, _, ids_batch) in dataset.iterbatches(

        batch_size, deterministic=True):

      n_samples = len(X_batch)

      y_pred_batch = self.predict_on_batch(X_batch)

      # Discard any padded predictions

      y_pred_batch = y_pred_batch[:n_samples]

      y_pred_batch = undo_transforms(y_pred_batch, transformers)

      y_preds.append(y_pred_batch)

    y_pred = np.concatenate(y_preds)

    return y_pred

  def evaluate(self, dataset, metrics, transformers=[], per_task_metrics=False):

    """

    Evaluates the performance of this model on specified dataset.

    Parameters

    ----------

    dataset: dc.data.Dataset

      Dataset object.

    metric: deepchem.metrics.Metric

      Evaluation metric

    transformers: list

      List of deepchem.transformers.Transformer

    per_task_metrics: bool

      If True, return per-task scores.

    Returns

    -------

    dict

      Maps tasks to scores under metric.

    """

    evaluator = Evaluator(self, dataset, transformers)

    if not per_task_metrics:

      scores = evaluator.compute_model_performance(metrics)

      return scores

    else:

      scores, per_task_scores = evaluator.compute_model_performance(

          metrics, per_task_metrics=per_task_metrics)

      return scores, per_task_scores