Update handler.py
Browse files- handler.py +3 -2
handler.py
CHANGED
@@ -5,7 +5,7 @@ import io
|
|
5 |
class EndpointHandler:
|
6 |
def __init__(self, model_dir=None):
|
7 |
# Set model size, assuming installation has been done with appropriate model files and setup
|
8 |
-
model_size = "
|
9 |
# Change to 'cuda' to use the GPU, and set compute_type for faster computation
|
10 |
self.model = WhisperModel(model_size, device="cuda", compute_type="float16")
|
11 |
|
@@ -20,7 +20,8 @@ class EndpointHandler:
|
|
20 |
segments, info = self.model.transcribe(audio_file)
|
21 |
|
22 |
# Compile the results into a text string and extract language information
|
23 |
-
|
|
|
24 |
language_code = info.language
|
25 |
language_prob = info.language_probability
|
26 |
|
|
|
5 |
class EndpointHandler:
|
6 |
def __init__(self, model_dir=None):
|
7 |
# Set model size, assuming installation has been done with appropriate model files and setup
|
8 |
+
model_size = "large-v2" if model_dir is None else model_dir
|
9 |
# Change to 'cuda' to use the GPU, and set compute_type for faster computation
|
10 |
self.model = WhisperModel(model_size, device="cuda", compute_type="float16")
|
11 |
|
|
|
20 |
segments, info = self.model.transcribe(audio_file)
|
21 |
|
22 |
# Compile the results into a text string and extract language information
|
23 |
+
# Strip whitespace from each segment before joining them
|
24 |
+
text = " ".join(segment.text.strip() for segment in segments)
|
25 |
language_code = info.language
|
26 |
language_prob = info.language_probability
|
27 |
|