Changed the parameters to be in env variables

handler.py

@@ -1,7 +1,8 @@
-from typing import Dict, List, Any, Literal
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
+from typing import Dict, List, Any, Literal, Optional, Tuple
 import torch
 import logging
+from pydantic_settings import BaseSettings
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -36,32 +37,38 @@ class EndpointHandler():
         """
         # get inputs
         inputs = data.pop("inputs", data)
-        logging.info(data.get("parameters"))
-        parameters = WhisperKwargsHandler(data.get("parameters", {}))  # Changed from pop to get
+        whisper_parameter_handler = WhisperParameterHandler()
+        logging.info(whisper_parameter_handler.model_dump())
 
 
         # run normal prediction
         prediction = self.pipe(
             inputs, 
-            **parameters.to_kwargs()
+            **whisper_parameter_handler.to_kwargs()
             )
         return prediction
     
-class WhisperKwargsHandler:
-    def __init__(self, kwargs: dict):
-        self.language: str = kwargs.get('language')
-        self.max_new_tokens: int = kwargs.get('max_new_tokens')
-        self.num_beams: int = kwargs.get('num_beams')
-        self.condition_on_prev_tokens: bool = kwargs.get('condition_on_prev_tokens')
-        self.compression_ratio_threshold: float = kwargs.get('compression_ratio_threshold')
-        self.temperature: tuple[float] = kwargs.get('temperature')
-        self.logprob_threshold: float = kwargs.get('logprob_threshold')
-        self.no_speech_threshold: float = kwargs.get('no_speech_threshold')
-        self.return_timestamps: Literal["word", True] = kwargs.get('return_timestamps')
-    
+
+class WhisperParameterHandler(BaseSettings):
+    language: Optional[str] = None  # Optional fields default to None
+    max_new_tokens: Optional[int] = None
+    num_beams: Optional[int] = None
+    condition_on_prev_tokens: Optional[bool] = None
+    compression_ratio_threshold: Optional[float] = None
+    temperature: Optional[Tuple[float, ...]] = None  # Optional Tuple
+    logprob_threshold: Optional[float] = None
+    no_speech_threshold: Optional[float] = None
+    return_timestamps: Optional[Literal["word", True]] = None
+
+    model_config = {
+        "env_prefix": "WHISPER_KWARGS_",
+        "case_sensitive": False,
+    }
+
     def to_kwargs(self):
-        """Convert object attributes to kwargs dict, excluding None values"""
+        """Convert object attributes to kwargs dict, excluding None values."""
         return {
-            key: value for key, value in self.__dict__.items() 
-            if value is not None
+            key: value
+            for key, value in self.model_dump().items()  # Use model_dump for accurate representation
+            if value is not None  # Exclude None values
         }

requirements.txt

@@ -1 +1,3 @@
-torch
+torch
+pydantic-settings
+transformers