From 89f541e988690632f24554b915bc171f4c276e57 Mon Sep 17 00:00:00 2001
From: Matias Fanger <maattz@outlook.com>
Date: Fri, 31 Mar 2023 20:22:09 -0300
Subject: [PATCH] fix: typos & weird extra lines

---
 common/batch.py                             | 3 +--
 common/checkpointing/snapshot.py            | 3 ---
 common/log_weights.py                       | 1 -
 core/custom_training_loop.py                | 1 -
 core/metric_mixin.py                        | 2 --
 core/metrics.py                             | 1 -
 core/train_pipeline.py                      | 2 --
 machines/list_ops.py                        | 1 -
 metrics/aggregation.py                      | 2 --
 metrics/auroc.py                            | 1 -
 metrics/rce.py                              | 3 ---
 ml_logging/absl_logging.py                  | 1 -
 ml_logging/torch_logging.py                 | 1 -
 model.py                                    | 1 -
 projects/home/recap/data/dataset.py         | 1 -
 projects/home/recap/data/util.py            | 2 --
 projects/home/recap/model/entrypoint.py     | 2 +-
 projects/home/recap/model/model_and_loss.py | 1 -
 projects/home/recap/optimizer/optimizer.py  | 1 -
 projects/twhin/models/models.py             | 1 -
 reader/dataset.py                           | 1 -
 reader/dds.py                               | 4 +---
 22 files changed, 3 insertions(+), 33 deletions(-)

diff --git a/common/batch.py b/common/batch.py
index cfd0ac0..4f6154d 100644
--- a/common/batch.py
+++ b/common/batch.py
@@ -1,5 +1,4 @@
-"""Extension of torchrec.dataset.utils.Batch to cover any dataset.
-"""
+"""Extension of torchrec.dataset.utils.Batch to cover any dataset."""
 # flake8: noqa
 from __future__ import annotations
 from typing import Dict
diff --git a/common/checkpointing/snapshot.py b/common/checkpointing/snapshot.py
index 2703efd..ff0f5ac 100644
--- a/common/checkpointing/snapshot.py
+++ b/common/checkpointing/snapshot.py
@@ -16,7 +16,6 @@ class Snapshot:
   """Checkpoints using torchsnapshot.
 
   Also saves step to be updated by the training loop.
-
   """
 
   def __init__(self, save_dir: str, state: Dict[str, Any]) -> None:
@@ -146,7 +145,6 @@ def checkpoints_iterator(save_dir: str, seconds_to_sleep: int = 30, timeout: int
   Args:
     seconds_to_sleep: time between polling calls.
     timeout: how long to wait for a new checkpoint.
-
   """
 
   def _poll(last_checkpoint: Optional[str] = None):
@@ -184,7 +182,6 @@ def get_checkpoint(
   Args:
     global_step: Finds this checkpoint if specified.
     missing_ok: if True and checkpoints do not exist, returns without restoration.
-
   """
   checkpoints = get_checkpoints(save_dir)
   if not checkpoints:
diff --git a/common/log_weights.py b/common/log_weights.py
index d2c8e21..763b6a7 100644
--- a/common/log_weights.py
+++ b/common/log_weights.py
@@ -18,7 +18,6 @@ def weights_to_log(
     model: model to traverse.
     how_to_log: if a function, then applies this to every parameter, if a dict
       then only applies and logs specified parameters.
-
   """
   if not how_to_log:
     return
diff --git a/core/custom_training_loop.py b/core/custom_training_loop.py
index 0241145..b05024f 100644
--- a/core/custom_training_loop.py
+++ b/core/custom_training_loop.py
@@ -6,7 +6,6 @@ Features (go/100_enablements):
     - On/off-chief evaluation
     - Warmstart/checkpoint management
     - go/dataset-service 0-copy integration
-
 """
 import datetime
 import os
diff --git a/core/metric_mixin.py b/core/metric_mixin.py
index a716ca7..2dfe727 100644
--- a/core/metric_mixin.py
+++ b/core/metric_mixin.py
@@ -22,7 +22,6 @@ for its update method.
 2. Redefine an existing metric class.
 
   SumMetric = prepend_transform(SumMetric, lambda outputs: {'value': 1})
-
 """
 from abc import abstractmethod
 from typing import Callable, Dict, List
@@ -88,7 +87,6 @@ def prepend_transform(base_metric: torchmetrics.Metric, transform: Callable):
 
   Functionally the same using inheritance, just saves some lines of code
   if no need for class attributes.
-
   """
 
   def transform_method(_self, *args, **kwargs):
diff --git a/core/metrics.py b/core/metrics.py
index 2384e4d..04102ee 100644
--- a/core/metrics.py
+++ b/core/metrics.py
@@ -1,7 +1,6 @@
 """Common metrics that also support multi task.
 
 We assume multi task models will output [task_idx, ...] predictions
-
 """
 from typing import Any, Dict
 
diff --git a/core/train_pipeline.py b/core/train_pipeline.py
index cde587e..6351d87 100644
--- a/core/train_pipeline.py
+++ b/core/train_pipeline.py
@@ -1,7 +1,6 @@
 """
 Taken from https://raw.githubusercontent.com/pytorch/torchrec/v0.3.2/torchrec/distributed/train_pipeline.py
 with TrainPipelineSparseDist.progress modified to support gradient accumulation.
-
 """
 import abc
 from dataclasses import dataclass, field
@@ -529,7 +528,6 @@ class TrainPipelineSparseDist(TrainPipeline[In, Out]):
     optimizer update/reset is skipped for `_grad_accum` calls of `progress`
     (congruent to training steps), and then update/reset on every `_grad_accum`th
     step.
-
     """
     should_step_optimizer = (
       self._grad_accum is not None
diff --git a/machines/list_ops.py b/machines/list_ops.py
index 71d5201..1fcfe57 100644
--- a/machines/list_ops.py
+++ b/machines/list_ops.py
@@ -15,7 +15,6 @@ Args:
 Typical usage would be in a bash script, e.g.:
 
   LIST_LEN=$(python list_ops.py --input_list=$INPUT --op=len)
-
 """
 import tml.machines.environment as env
 
diff --git a/metrics/aggregation.py b/metrics/aggregation.py
index 57f75d5..4c68636 100644
--- a/metrics/aggregation.py
+++ b/metrics/aggregation.py
@@ -24,7 +24,6 @@ def update_mean(
     weight: The weights for the new value.
 
   Returns: The updated mean and updated weighted sum.
-
   """
   weight = torch.broadcast_to(weight, value.shape)
 
@@ -42,7 +41,6 @@ def stable_mean_dist_reduce_fn(state: torch.Tensor) -> torch.Tensor:
     state: A tensor with the first dimension indicating workers.
 
   Returns: The accumulated mean from all workers.
-
   """
   mean, weight_sum = update_mean(
     current_mean=torch.as_tensor(0.0, dtype=state.dtype, device=state.device),
diff --git a/metrics/auroc.py b/metrics/auroc.py
index 6979c20..39bc3db 100644
--- a/metrics/auroc.py
+++ b/metrics/auroc.py
@@ -62,7 +62,6 @@ class AUROCWithMWU(torchmetrics.Metric):
 
   def __init__(self, label_threshold: float = 0.5, raise_missing_class: bool = False, **kwargs):
     """
-
     Args:
       label_threshold: Labels strictly above this threshold are considered positive labels,
                        otherwise, they are considered negative.
diff --git a/metrics/rce.py b/metrics/rce.py
index b6ada1d..efd6cef 100644
--- a/metrics/rce.py
+++ b/metrics/rce.py
@@ -41,7 +41,6 @@ def _binary_cross_entropy_with_clipping(
     reduction: The reduction method to use.
 
   Returns: Binary cross entropy on the clipped predictions.
-
   """
   predictions = torch.clamp(predictions, epsilon, 1.0 - epsilon)
   bce = -target * torch.log(predictions + epsilon)
@@ -220,14 +219,12 @@ class NRCE(RCE):
             perform if it was well calibrated.
 
   .. note:: A big gap between NRCE and RCE might indicate a badly calibrated model,
-
   """
 
   def __init__(
     self, from_logits: bool = False, label_smoothing: float = 0, epsilon: float = 1e-7, **kwargs
   ):
     """
-
     Args:
       from_logits: whether or not predictions are logits or probabilities.
       label_smoothing: label smoothing constant.
diff --git a/ml_logging/absl_logging.py b/ml_logging/absl_logging.py
index e4d9e9e..47a8686 100644
--- a/ml_logging/absl_logging.py
+++ b/ml_logging/absl_logging.py
@@ -5,7 +5,6 @@
 Usage:
     >>> from twitter.ml.logging.absl_logging import logging
     >>> logging.info(f"Properly logged as INFO level in GCP Stackdriver.")
-
 """
 import logging as py_logging
 import sys
diff --git a/ml_logging/torch_logging.py b/ml_logging/torch_logging.py
index e791c46..ff1b2f4 100644
--- a/ml_logging/torch_logging.py
+++ b/ml_logging/torch_logging.py
@@ -6,7 +6,6 @@
     >>> from ml.logging.torch_logging import logging
     >>> logging.info(f"This only prints on rank 0 if distributed, otherwise prints normally.")
     >>> logging.info(f"This prints on all ranks if distributed, otherwise prints normally.", rank=-1)
-
 """
 import functools
 from typing import Optional
diff --git a/model.py b/model.py
index 9df13bc..e53bc87 100644
--- a/model.py
+++ b/model.py
@@ -33,7 +33,6 @@ class ModelAndLoss(torch.nn.Module):
     prefetching purposes during training using torchrec's pipeline.  However
     the underlying model signature needs to be exportable to onnx, requiring
     generic python types.  see https://pytorch.org/docs/stable/onnx.html#types.
-
     """
     outputs = self.model(batch)
     losses = self.loss_fn(outputs["logits"], batch.labels.float(), batch.weights.float())
diff --git a/projects/home/recap/data/dataset.py b/projects/home/recap/data/dataset.py
index 3478c68..228eb82 100644
--- a/projects/home/recap/data/dataset.py
+++ b/projects/home/recap/data/dataset.py
@@ -155,7 +155,6 @@ def get_explicit_datetime_inputs_files(explicit_datetime_inputs):
   Returns:
     data_files: Sorted list of files to read corresponding to data at the desired datetimes
     num_hours_missing: Number of hours that we are missing data
-
   """
   datetimes = get_datetimes(explicit_datetime_inputs)
   folders = [os.path.join(explicit_datetime_inputs.data_root, datetime) for datetime in datetimes]
diff --git a/projects/home/recap/data/util.py b/projects/home/recap/data/util.py
index a9fd51e..ad3d6d0 100644
--- a/projects/home/recap/data/util.py
+++ b/projects/home/recap/data/util.py
@@ -14,7 +14,6 @@ def keyed_tensor_from_tensors_dict(
     tensor_map:
 
   Returns:
-
   """
   keys = list(tensor_map.keys())
   # We expect batch size to be first dim. However, if we get a shape [Batch_size],
@@ -64,7 +63,6 @@ def keyed_jagged_tensor_from_tensors_dict(
     tensor_map:
 
   Returns:
-
   """
 
   if not tensor_map:
diff --git a/projects/home/recap/model/entrypoint.py b/projects/home/recap/model/entrypoint.py
index 8f4d534..c76219d 100644
--- a/projects/home/recap/model/entrypoint.py
+++ b/projects/home/recap/model/entrypoint.py
@@ -26,7 +26,7 @@ def unsanitize(sanitized_task_name):
 
 
 def _build_single_task_model(task: model_config_mod.TaskModel, input_shape: int):
-  """ "Builds a model for a single task"""
+  """"Builds a model for a single task"""
   if task.mlp_config:
     return mlp.Mlp(in_features=input_shape, mlp_config=task.mlp_config)
   elif task.dcn_config:
diff --git a/projects/home/recap/model/model_and_loss.py b/projects/home/recap/model/model_and_loss.py
index b42da9b..54c5b7a 100644
--- a/projects/home/recap/model/model_and_loss.py
+++ b/projects/home/recap/model/model_and_loss.py
@@ -29,7 +29,6 @@ class ModelAndLoss(torch.nn.Module):
     prefetching purposes during training using torchrec's pipeline.  However
     the underlying model signature needs to be exportable to onnx, requiring
     generic python types.  see https://pytorch.org/docs/stable/onnx.html#types.
-
     """
     outputs = self.model(
       continuous_features=batch.continuous_features,
diff --git a/projects/home/recap/optimizer/optimizer.py b/projects/home/recap/optimizer/optimizer.py
index c5b0cf1..dd08b35 100644
--- a/projects/home/recap/optimizer/optimizer.py
+++ b/projects/home/recap/optimizer/optimizer.py
@@ -27,7 +27,6 @@ class RecapLRShim(torch.optim.lr_scheduler._LRScheduler):
 
   This adheres to the torch.optim scheduler API and can be plugged anywhere that
   e.g. exponential decay can be used.
-
   """
 
   def __init__(
diff --git a/projects/twhin/models/models.py b/projects/twhin/models/models.py
index b1dfaa7..6a2c7c9 100644
--- a/projects/twhin/models/models.py
+++ b/projects/twhin/models/models.py
@@ -142,7 +142,6 @@ class TwhinModelAndLoss(torch.nn.Module):
     prefetching purposes during training using torchrec's pipeline.  However
     the underlying model signature needs to be exportable to onnx, requiring
     generic python types.  see https://pytorch.org/docs/stable/onnx.html#types.
-
     """
     outputs = self.model(batch)
     logits = outputs["logits"]
diff --git a/reader/dataset.py b/reader/dataset.py
index 6e811cc..2a30ca2 100644
--- a/reader/dataset.py
+++ b/reader/dataset.py
@@ -2,7 +2,6 @@
 
 - Override `pa_to_batch` for dataset specific imputation, negative sampling, or coercion to Batch.
 - Readers can be colocated or off trainer machines.
-
 """
 import abc
 import functools
diff --git a/reader/dds.py b/reader/dds.py
index 7b49893..d70331d 100644
--- a/reader/dds.py
+++ b/reader/dds.py
@@ -1,5 +1,4 @@
-"""Dataset service orchestrated by a TFJob
-"""
+"""Dataset service orchestrated by a TFJob"""
 from typing import Optional
 import uuid
 
@@ -104,7 +103,6 @@ def maybe_distribute_dataset(dataset: tf.data.Dataset) -> tf.data.Dataset:
   - all rank processes will consume from the same job/dataset.
 
   Without this, dataset workers will try to serve 1 job per rank process and OOM.
-
   """
   if not env.has_readers():
     return dataset