| from typing import Dict, List, Any |
| from transformers import pipeline |
|
|
| import sys |
| import torch |
| from transformers import ( |
| AutomaticSpeechRecognitionPipeline, |
| WhisperForConditionalGeneration, |
| WhisperTokenizer, |
| WhisperProcessor |
| ) |
| from peft import LoraConfig, PeftModel, LoraModel, LoraConfig, get_peft_model, PeftConfig |
|
|
| class EndpointHandler(): |
| def __init__(self, path=""): |
| |
| language = "Chinese" |
| task = "transcribe" |
| peft_config = PeftConfig.from_pretrained(path) |
| model = WhisperForConditionalGeneration.from_pretrained( |
| peft_config.base_model_name_or_path |
| ) |
| 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 = 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`) |
| date (:obj: `str`) |
| Return: |
| A :obj:`list` | `dict`: will be serialized and returned |
| """ |
| |
| |
| |
| 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 |
