Update pipeline.py

pipeline.py

@@ -280,19 +280,28 @@ def binary_search_above(transcript, start_index, end_index, time):
 
 class PreTrainedPipeline():
     def __init__(self, path: str):
-        path2 = os.path.join(path, "model")
+        path2 = os.path.join(path, 'model')
         self.model2 = AutoModelForSequenceClassification.from_pretrained(path2)
         self.tokenizer2 = AutoTokenizer.from_pretrained(path2)
         self.pipeline2 = SponsorBlockClassificationPipeline(
             model=self.model2, tokenizer=self.tokenizer2)
     
     def __call__(self, inputs: str)-> List[Dict[str, Any]]:
-        split = inputs.split(',')
-        data = {
-            'video_id': split[0],
-            'start': float(split[1]),
-            'end': float(split[2])
-        }
+        
+
+        if ' ' not in inputs and inputs.count(',') >= 2: # Automated call (compressed string)
+            split_info = inputs.split(',', 1)
+            times = np.reshape(np.array(split_info[1].split(',')), (-1, 2))
+            data = []
+            for start, end in times:
+                data.append({
+                    'video_id': split_info[0],
+                    'start': float(start),
+                    'end': float(end)
+                })
+        else:
+            data = inputs
+
         return self.pipeline2(data)
 
 
@@ -301,26 +310,22 @@ class SponsorBlockClassificationPipeline(TextClassificationPipeline):
     def __init__(self, model, tokenizer):
         super().__init__(model=model, tokenizer=tokenizer, return_all_scores=True)
 
-    def preprocess(self, video, **tokenizer_kwargs):
+    def preprocess(self, data, **tokenizer_kwargs):
+        if isinstance(data, str): # If string, assume this is what user wants to classify
+            text = data 
+        else: # Otherwise, get data from transcript
+            words = get_words(video['video_id'])
+            segment_words = extract_segment(words, video['start'], video['end'])
+            text = ' '.join(x['text'] for x in segment_words)
 
-        words = get_words(video['video_id'])
-        segment_words = extract_segment(words, video['start'], video['end'])
-        text = ' '.join(x['text'] for x in segment_words)
-
-        model_inputs = self.tokenizer(
+        return self.tokenizer(
             text, return_tensors=self.framework, **tokenizer_kwargs)
-        return {'video': video, 'model_inputs': model_inputs}
-
-    def _forward(self, data):
-        model_outputs = self.model(**data['model_inputs'])
-        return {'video': data['video'], 'model_outputs': model_outputs}
 
-    def postprocess(self, data, function_to_apply=None, return_all_scores=False):
-        model_outputs = data['model_outputs']
 
+    def postprocess(self, model_outputs, function_to_apply=None, return_all_scores=False):
         results = super().postprocess(model_outputs, function_to_apply, return_all_scores)
 
         for result in results:
             result['label_text'] = CATEGORIES[result['label']]
 
-        return results # {**data['video'], 'result': results}
+        return results