|
|
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 |
