modify handler

handler.py

@@ -24,8 +24,8 @@ class EndpointHandler():
         processor = WhisperProcessor.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
         feature_extractor = processor.feature_extractor
         self.forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
-        # self.pipeline = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
-        self.pipeline = pipeline(task= "automatic-speech-recognition", model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
+        self.pipeline = AutomaticSpeechRecognitionPipeline(model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
+        # self.pipeline = pipeline(task= "automatic-speech-recognition", model=model, tokenizer=tokenizer, feature_extractor=feature_extractor)
         self.pipeline.model.config.forced_decoder_ids = self.pipeline.tokenizer.get_decoder_prompt_ids(language="Chinese", task="transcribe")
         self.pipeline.model.generation_config.forced_decoder_ids = self.pipeline.model.config.forced_decoder_ids  # just to be sure!
         # self.pipeline = pipeline(task= "automatic-speech-recognition", model=self.model)
@@ -43,7 +43,7 @@ class EndpointHandler():
 
         inputs = data.pop("inputs", data)
         with torch.cuda.amp.autocast():
-            # prediction = self.pipeline(inputs, generate_kwargs={"forced_decoder_ids": self.forced_decoder_ids}, max_new_tokens=255)["text"]
-            prediction = self.pipeline(inputs, return_timestamps=False)
+            prediction = self.pipeline(inputs, generate_kwargs={"forced_decoder_ids": self.forced_decoder_ids}, max_new_tokens=255)
+            # prediction = self.pipeline(inputs, return_timestamps=False)
             prediction['text'] = prediction['text'] + '????'
         return prediction