update handler

handler.py

@@ -1,48 +1,53 @@
-import torch
 from typing import Dict, List, Any
+from transformers import pipeline
+
+import sys
+import torch
 from transformers import (
-    AutomaticSpeechRecognitionPipeline,
-    WhisperForConditionalGeneration,
-    WhisperTokenizer,
-    WhisperProcessor,
-    pipeline
+AutomaticSpeechRecognitionPipeline,
+WhisperForConditionalGeneration,
+WhisperTokenizer,
+WhisperProcessor
 )
 from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model, PeftConfig
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # Preload all the elements you are going to need at inference.
-        peft_model_id = "cathyi/openai-whisper-large-v2-Lora"
+        
         language = "Chinese"
-        task = "transcribe" 
-        peft_config = PeftConfig.from_pretrained(peft_model_id)
-        model= WhisperForConditionalGeneration.from_pretrained(
+        task = "transcribe"
+        peft_config = PeftConfig.from_pretrained(path)
+        model = WhisperForConditionalGeneration.from_pretrained(
             peft_config.base_model_name_or_path
         )
-        model = PeftModel.from_pretrained(model, peft_model_id)
+        model = PeftModel.from_pretrained(model, path)
         tokenizer = WhisperTokenizer.from_pretrained(peft_config.base_model_name_or_path, language=language, task=task)
         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.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)
-        # 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
-
+        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=language, task=task)
+        self.pipeline.model.generation_config.forced_decoder_ids = self.pipeline.model.config.forced_decoder_ids
+    
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-       data args:
-            inputs (:obj: `str` | `PIL.Image` | `np.array`)
-            kwargs
-      Return:
+        data args:
+            inputs (:obj: `str`)
+            date (:obj: `str`)
+        Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-
+        # get inputs
+        
+        # run normal prediction
         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)
+        print("a1", inputs)
+        print("a2", inputs, file=sys.stderr)
+        print("a3", inputs, file=sys.stdout)
+
+        prediction = self.pipeline(inputs, return_timestamps=False)
+
+        print("b1", prediction)
+        print("b2", prediction, file=sys.stderr)
+        print("b3", prediction, file=sys.stdout)
         return prediction