From de2508867d0bb09b449010899c027230880ec558 Mon Sep 17 00:00:00 2001 From: shiyu22 Date: Wed, 31 May 2023 14:58:04 +0800 Subject: [PATCH] Update pipelines with pydantic Signed-off-by: shiyu22 --- audio_embedding.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/audio_embedding.py b/audio_embedding.py index aaf1f33..7edbc94 100644 --- a/audio_embedding.py +++ b/audio_embedding.py @@ -1,20 +1,22 @@ +from posixpath import basename +from typing import Optional +from pydantic import BaseModel from towhee import pipe, ops, AutoPipes, AutoConfig @AutoConfig.register -class AudioEmbeddingConfig: - def __init__(self): - # config for audio_decode.ffmpeg - self.batch_size = -1 - self.sample_rate = None - self.layout = None - - # config for audio_embedding.vggish - self.weights_path: str = None - self.framework: str = 'pytorch' - - # config for triton - self.device = -1 +class AudioEmbeddingConfig(BaseModel): + # config for audio_decode.ffmpeg + batch_size: Optional[int] = -1 + sample_rate: Optional[float] = None + layout: Optional[str] = None + + # config for audio_embedding.vggish + weights_path: Optional[str] = None + framework: Optional[str] = 'pytorch' + + # config for triton + device: Optional[int] = -1 @AutoPipes.register