Fix linting issues

Author: joecummings

test/torchtext_unittest/prototype/test_generate.py

@@ -1,8 +1,9 @@
 from unittest.mock import patch
+
+import torch
 from torchtext.prototype.generate import GenerationUtil
 from torchtext.prototype.models import T5_BASE_GENERATION
 from torchtext_unittest.common.torchtext_test_case import TorchtextTestCase
-import torch
 
 
 class TestGenerationUtil(TorchtextTestCase):

torchtext/prototype/generate.py

@@ -1,11 +1,10 @@
+import logging
 from typing import Optional
 
 import torch
 import torch.nn.functional as F
 from torch import nn
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 

torchtext/prototype/models/t5/__init__.py

@@ -1,19 +1,19 @@
 from .bundler import (
-    T5_BASE_ENCODER,
+    T5_11B,
+    T5_11B_ENCODER,
+    T5_11B_GENERATION,
+    T5_3B,
+    T5_3B_ENCODER,
+    T5_3B_GENERATION,
     T5_BASE,
+    T5_BASE_ENCODER,
     T5_BASE_GENERATION,
-    T5_SMALL_ENCODER,
-    T5_SMALL,
-    T5_SMALL_GENERATION,
-    T5_LARGE_ENCODER,
     T5_LARGE,
+    T5_LARGE_ENCODER,
     T5_LARGE_GENERATION,
-    T5_3B_ENCODER,
-    T5_3B,
-    T5_3B_GENERATION,
-    T5_11B_ENCODER,
-    T5_11B,
-    T5_11B_GENERATION,
+    T5_SMALL,
+    T5_SMALL_ENCODER,
+    T5_SMALL_GENERATION,
     T5Bundle,
 )
 from .model import T5Conf, T5Model

torchtext/prototype/models/t5/model.py

@@ -1,14 +1,13 @@
+# logging library is not automatically supported by Torchscript
+import warnings
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Union, Callable
+from typing import Callable, Dict, List, Optional, Union
 
 import torch
 import torch.nn as nn
 from torch import Tensor
 
-from .modules import T5Encoder, T5Decoder
-
-# logging library is not automatically supported by Torchscript
-import warnings
+from .modules import T5Decoder, T5Encoder
 
 
 @dataclass(frozen=True)
@@ -194,7 +193,6 @@ def forward(
             )
 
         if not self.encoder_only:
-
             assert self.decoder is not None
             assert encoder_outputs is not None
 
@@ -238,7 +236,7 @@ def forward(
                 # Rescale output before projecting on vocab. This happens when the encoder and decoder share the
                 # same word embeddings, which is always the case in our t5 implementation.
                 # See https://github.com/huggingface/transformers/blob/d0acc9537829e7d067edbb791473bbceb2ecf056/src/transformers/models/t5/modeling_t5.py#L1661
-                decoder_output = decoder_output * (self.embedding_dim**-0.5)
+                decoder_output = decoder_output * (self.embedding_dim ** -0.5)
                 decoder_output = self.lm_head(decoder_output)
                 decoder_outputs["decoder_output"] = decoder_output
 

torchtext/prototype/models/t5/modules.py

@@ -15,7 +15,7 @@
 
 import math
 import warnings
-from typing import Dict, List, Optional, Tuple, Union, Callable
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -843,7 +843,7 @@ def forward(
         embedded_tgt: Optional[Tensor] = None,
     ) -> Dict[str, Union[Tensor, List[Tensor], Optional[Tensor], List[Optional[Tensor]]]]:
         r"""Pass the input (and masks) through the stack of encoder layers.
-        
+
         Args:
             tgt (Optional[Tensor]): Tokenized input sequence to the encoder.
                 Must be batch first with shape (B, Ne) where B is the batch size and Ne is the
@@ -857,7 +857,7 @@ def forward(
                 length, and E is the model dimension.
                 *Note*: If you do not provide this `embedded_tgt`, you must have provided a `token_embedding` layer \
                     in the initialization of the T5Encoder.
-        
+
         Returns:
             Tuple of last hidden layer, all hidden layers, position bias, and self-attention scores
         """

torchtext/prototype/models/t5/wrapper.py

@@ -1,22 +1,21 @@
+import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from torch import Tensor
 from torchtext.prototype.models import (
+    T5_11B_GENERATION,
+    T5_3B_GENERATION,
     T5_BASE_GENERATION,
-    T5_SMALL_GENERATION,
     T5_LARGE_GENERATION,
-    T5_3B_GENERATION,
-    T5_11B_GENERATION,
+    T5_SMALL_GENERATION,
+    T5Bundle,
     T5Conf,
     T5Transform,
-    T5Bundle,
 )
 
-import warnings
-
 
 BUNDLERS = {
     "base": T5_BASE_GENERATION,
@@ -139,7 +138,6 @@ def beam_search(
         return new_decoder_tokens, new_scores, new_incomplete_sentences
 
     def generate(self, encoder_tokens: Tensor, beam_size: int, eos_idx: int = 1, max_seq_len: int = 512) -> Tensor:
-
         # pass tokens through encoder
         bsz = encoder_tokens.size(0)
         encoder = self.model.get_encoder()
@@ -155,7 +153,6 @@ def generate(self, encoder_tokens: Tensor, beam_size: int, eos_idx: int = 1, max
 
         # iteratively generate output sequence until all sequences in the batch have generated the end-of-sequence token
         for step in range(max_seq_len):
-
             if step == 1:
                 # duplicate and order encoder output so that each beam is treated as its own independent sequence
                 encoder_output = encoder_outputs.get("encoder_output")
@@ -189,7 +186,6 @@ def generate(self, encoder_tokens: Tensor, beam_size: int, eos_idx: int = 1, max
         return decoder_tokens
 
     def forward(self, input_text: List[str], beam_size: int, max_seq_len: int) -> Union[List[str], str]:
-
         model_input = self.transform(input_text)
         model_output_tensor = self.generate(encoder_tokens=model_input, beam_size=beam_size, max_seq_len=max_seq_len)
         model_output_list = torch.jit.annotate(List[List[int]], model_output_tensor.tolist())