update handler

.gitignore

@@ -0,0 +1 @@
+test_handler.py

handler.py

@@ -1,3 +1,4 @@
+import torch
 from typing import Dict, List, Any
 from transformers import (
     AutomaticSpeechRecognitionPipeline,
@@ -10,9 +11,10 @@ from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model, P
 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" 
-        self.peft_config = PeftConfig.from_pretrained(path)
+        peft_config = PeftConfig.from_pretrained(peft_model_id)
         model= WhisperForConditionalGeneration.from_pretrained(
             peft_config.base_model_name_or_path
         )
@@ -20,9 +22,8 @@ class EndpointHandler():
         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
-        forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task=task)
+        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=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
@@ -38,5 +39,5 @@ class EndpointHandler():
 
         inputs = data.pop("inputs", data)
         with torch.cuda.amp.autocast():
-            predicion = self.pipeline(inputs, generate_kwargs={"forced_decoder_ids": forced_decoder_ids}, max_new_tokens=255)["text"]
+            prediction = self.pipeline(inputs, generate_kwargs={"forced_decoder_ids": self.forced_decoder_ids}, max_new_tokens=255)["text"]
         return {"prediction": prediction}