Update pipeline.py

pipeline.py

@@ -6,14 +6,6 @@ from accelerate.utils import get_balanced_memory
 from transformers import BitsAndBytesConfig, StoppingCriteria, StoppingCriteriaList
 from typing import Dict, List, Any
 
-# define custom stopping criteria object
-class StopOnTokens(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        for stop_ids in stop_token_ids:
-            if torch.eq(input_ids[0][-len(stop_ids):], stop_ids).all():
-                return True
-        return False
-
 class PreTrainedPipeline():
     def __init__(self, path=""):
         path = "oleksandrfluxon/mpt-7b-instruct-evaluate"
@@ -44,6 +36,14 @@ class PreTrainedPipeline():
         stop_token_ids = [torch.LongTensor(x).to(device) for x in stop_token_ids]
         print("===> stop_token_ids", stop_token_ids)
 
+        # define custom stopping criteria object
+        class StopOnTokens(StoppingCriteria):
+            def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+                for stop_ids in stop_token_ids:
+                    if torch.eq(input_ids[0][-len(stop_ids):], stop_ids).all():
+                        return True
+                return False
+
         stopping_criteria = StoppingCriteriaList([StopOnTokens()])
 
         self.pipeline = transformers.pipeline(