refactor(model): minor code style updates

basaran/model.py

@@ -117,6 +117,14 @@ def __call__(
                     **samples,
                 )
 
+    @retry(stop=stop_after_attempt(5), wait=wait_fixed(1))
+    def _infer(self, model_fn, **kwargs):
+        """Call a model function in inference mode with auto retrying."""
+        # This is a temporary workaround for bitsandbytes #162:
+        # https://github.com/TimDettmers/bitsandbytes/issues/162
+        with torch.inference_mode():
+            return model_fn(**kwargs)
+
     def _sample(self, token, token_logprob, top_tokens, top_logprobs):
         """Sample log probabilities of the most likely tokens."""
         token = self.tokenizer.decode(token)
@@ -168,14 +176,6 @@ def _logits_processor(self, config, input_length):
 
         return processor
 
-    @retry(stop=stop_after_attempt(5), wait=wait_fixed(1))
-    def _infer(self, model_fn, **kwargs):
-        """Call a model function in inference mode with auto retrying."""
-        # This is a temporary workaround for bitsandbytes #162:
-        # https://github.com/TimDettmers/bitsandbytes/issues/162
-        with torch.inference_mode():
-            return model_fn(**kwargs)
-
     def tokenize(self, text):
         """Tokenize a string into a tensor of token IDs."""
         batch = self.tokenizer.encode(text, return_tensors="pt")
@@ -280,7 +280,7 @@ def generate(self, input_ids, logprobs=0, **kwargs):
             # Append selected tokens to the inputs.
             input_ids = torch.cat([input_ids, tokens[:, None]], dim=-1)
 
-            # Extract past key values from model output.
+            # Extract past key values from model outputs.
             if "past_key_values" in outputs:
                 kwargs["past_key_values"] = outputs.past_key_values
             elif "mems" in outputs: