From 8ca47d95bd356d9b6d06b50b6e21501645e91cf6 Mon Sep 17 00:00:00 2001
From: evelynmitchell <efmphone@gmail.com>
Date: Wed, 3 Jan 2024 07:57:14 -0700
Subject: [PATCH] flake8 cleanup

---
 swarms/tools/logits_processor.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/swarms/tools/logits_processor.py b/swarms/tools/logits_processor.py
index ed7fef18..c6ba1691 100644
--- 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,
         )