Files changed (1) hide show
  1. handler.py +5 -3
handler.py CHANGED
@@ -8,16 +8,17 @@ class EndpointHandler():
8
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
9
  torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
10
  model_id = "openai/whisper-large-v3"
11
- self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
 
12
  model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
13
  )
14
- self.model.to(device)
15
 
16
  processor = AutoProcessor.from_pretrained(model_id)
17
 
18
  self.pipeline = pipeline(
19
  "automatic-speech-recognition",
20
- model=self.model,
21
  tokenizer=processor.tokenizer,
22
  feature_extractor=processor.feature_extractor,
23
  max_new_tokens=128,
@@ -27,6 +28,7 @@ class EndpointHandler():
27
  torch_dtype=torch_dtype,
28
  device=device,
29
  )
 
30
 
31
 
32
  def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]:
 
8
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
9
  torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
10
  model_id = "openai/whisper-large-v3"
11
+
12
+ model = AutoModelForSpeechSeq2Seq.from_pretrained(
13
  model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
14
  )
15
+ model.to(device)
16
 
17
  processor = AutoProcessor.from_pretrained(model_id)
18
 
19
  self.pipeline = pipeline(
20
  "automatic-speech-recognition",
21
+ model=model,
22
  tokenizer=processor.tokenizer,
23
  feature_extractor=processor.feature_extractor,
24
  max_new_tokens=128,
 
28
  torch_dtype=torch_dtype,
29
  device=device,
30
  )
31
+ self.model = model
32
 
33
 
34
  def __call__(self, data: Dict[str, bytes]) -> Dict[str, str]: