flake8 cleanup

1 year ago · 8ca47d95bd
parent 5b1c408376
commit 8ca47d95bd
1 changed files with 3 additions and 1 deletions
--- a/swarms/tools/logits_processor.py
+++ b/swarms/tools/logits_processor.py
@ -1,3 +1,5 @@
+"""Logits processors for the GPT-Neo model."""
+
 from transformers import (
    PreTrainedTokenizer,
    LogitsWarper,
@ -48,7 +50,7 @@ class NumberStoppingCriteria(StoppingCriteria):
        scores: torch.FloatTensor,
    ) -> bool:
        decoded = self.tokenizer.decode(
-            input_ids[0][self.prompt_length :],
+            input_ids[0][self.prompt_length:],
            skip_special_tokens=True,
        )