update handler for testing

handler.py

@@ -1,53 +1,48 @@
-from typing import Dict, List, Any
-from transformers import pipeline
-
-import sys
 import torch
+from typing import Dict, List, Any
 from transformers import (
-AutomaticSpeechRecognitionPipeline,
-WhisperForConditionalGeneration,
-WhisperTokenizer,
-WhisperProcessor
+    AutomaticSpeechRecognitionPipeline,
+    WhisperForConditionalGeneration,
+    WhisperTokenizer,
+    WhisperProcessor,
+    pipeline
 )
 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(path)
-        model = WhisperForConditionalGeneration.from_pretrained(
+        task = "transcribe" 
+        peft_config = PeftConfig.from_pretrained(peft_model_id)
+        model= WhisperForConditionalGeneration.from_pretrained(
             peft_config.base_model_name_or_path
         )
-        model = PeftModel.from_pretrained(model, path)
+        model = PeftModel.from_pretrained(model, peft_model_id)
         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 = 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
-    
+        # 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
+
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
-        data args:
-            inputs (:obj: `str`)
-            date (:obj: `str`)
-        Return:
+       data args:
+            inputs (:obj: `str` | `PIL.Image` | `np.array`)
+            kwargs
+      Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
-        # get inputs
-        
-        # run normal prediction
-        inputs = data.pop("inputs", data)
-        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
+        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)
+        return [prediction, {"test": 0}]