Skip to content

Commit 01fa12e

Browse files
fqscfqjfqscfqj
authored
feat(nemo): enable word-level timestamps for ASR models (#10297)
* feat(nemo): enable word-level timestamps for ASR models The nemo backend ignored timestamp_granularities and always returned a single segment with start=0 end=0, making word-level timestamps impossible to obtain even though the NeMo models (parakeet-tdt, etc.) fully support them. Changes: - Add _get_stride_seconds() to compute frame duration from the model's preprocessor window_stride and encoder subsampling_factor. - Add _build_segments_with_words() that extracts word offsets from the NeMo Hypothesis.timestamp dict and converts frame indices to nanosecond timestamps. - Support 'word' granularity (one segment per word) and 'segment' granularity (merge at time-gap boundaries using a dynamic threshold). - Populate TranscriptSegment.words with TranscriptWord entries so callers get both segment-level and word-level timing. - Only request timestamps from NeMo when the caller actually asks for them (timestamp_granularities is non-empty), keeping the fast path unchanged for callers that don't need timestamps. Tested with nvidia/parakeet-tdt-0.6b-v3 on the JFK "ask not" clip: curl -X POST /v1/audio/transcriptions \ -F file=@jfk.wav -F model=nemo-parakeet-tdt-0.6b \ -F 'timestamp_granularities[]=word' -F response_format=verbose_json → each word has correct start/end times in seconds. Signed-off-by: fqscfqj <fqscfqj@outlook.com> * fix(nemo): address Copilot review feedback - Narrow exception handling in _get_stride_seconds to catch only AttributeError, KeyError, TypeError instead of bare Exception, and emit a warning when falling back to the hardcoded stride. - Remove explicit return_hypotheses=False when timestamps are requested; timestamps=True already forces NeMo to return Hypothesis objects. - Add a warning when NeMo does not return Hypothesis objects despite timestamps being requested. Signed-off-by: fqscfqj <fqscfqj@outlook.com> --------- Signed-off-by: fqscfqj <fqscfqj@outlook.com>
1 parent cf7f957 commit 01fa12e

1 file changed

Lines changed: 186 additions & 16 deletions

File tree

‎backend/python/nemo/backend.py‎

Lines changed: 186 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,135 @@ def LoadModel(self, request, context):
8484

8585
return backend_pb2.Result(message="Model loaded successfully", success=True)
8686

87+
def _get_stride_seconds(self):
88+
"""Compute the seconds-per-frame stride for the loaded model.
89+
90+
stride = preprocessor_window_stride * encoder_subsampling_factor
91+
"""
92+
try:
93+
preprocessor = self.model.preprocessor
94+
window_stride = preprocessor._cfg.get('window_stride', 0.01)
95+
subsampling_factor = getattr(self.model.encoder, 'subsampling_factor', 8)
96+
return window_stride * subsampling_factor
97+
except (AttributeError, KeyError, TypeError) as err:
98+
print(
99+
f"Warning: could not compute stride from model config ({err}), "
100+
f"falling back to 0.08s/frame",
101+
file=sys.stderr,
102+
)
103+
return 0.08
104+
105+
def _build_segments_with_words(self, hypothesis, stride, timestamp_granularities=None):
106+
"""Build TranscriptSegment list from a NeMo Hypothesis with timestamps.
107+
108+
Supports two granularity modes:
109+
- "word": one TranscriptSegment per word, each with a single TranscriptWord entry
110+
- "segment" (default): merge consecutive words into sentence-level segments,
111+
splitting at word-level time gaps that exceed a dynamic threshold.
112+
"""
113+
if not hypothesis or not isinstance(hypothesis.timestamp, dict):
114+
return []
115+
116+
word_offsets = hypothesis.timestamp.get('word', [])
117+
if not word_offsets:
118+
return []
119+
120+
granularities = list(timestamp_granularities) if timestamp_granularities else []
121+
granularity = "word" if "word" in granularities else "segment"
122+
123+
# Build a flat list of (text, start_ns, end_ns) from NeMo word offsets
124+
transcript_words = []
125+
for wo in word_offsets:
126+
word_text = wo.get('word', '')
127+
if not word_text:
128+
continue
129+
start_offset = wo.get('start_offset', 0)
130+
end_offset = wo.get('end_offset', start_offset)
131+
start_ns = int(start_offset * stride * 1_000_000_000)
132+
end_ns = int(end_offset * stride * 1_000_000_000)
133+
transcript_words.append({
134+
'text': word_text,
135+
'start': start_ns,
136+
'end': end_ns,
137+
})
138+
139+
if not transcript_words:
140+
return []
141+
142+
if granularity == "word":
143+
# One segment per word
144+
result = []
145+
for idx, tw in enumerate(transcript_words):
146+
word = backend_pb2.TranscriptWord(
147+
start=tw['start'], end=tw['end'], text=tw['text']
148+
)
149+
result.append(backend_pb2.TranscriptSegment(
150+
id=idx,
151+
start=tw['start'],
152+
end=tw['end'],
153+
text=tw['text'],
154+
words=[word],
155+
))
156+
return result
157+
158+
# segment mode — merge at word-level time-gap boundaries
159+
# Compute gap threshold: median inter-word gap * 3, clamped to [0.3, 2.0]s
160+
gaps = []
161+
for i in range(1, len(transcript_words)):
162+
gap = (transcript_words[i]['start'] - transcript_words[i - 1]['end']) / 1_000_000_000
163+
if gap > 0:
164+
gaps.append(gap)
165+
if gaps:
166+
gaps.sort()
167+
median_gap = gaps[len(gaps) // 2]
168+
threshold_ns = int(max(0.3, min(median_gap * 3, 2.0)) * 1_000_000_000)
169+
else:
170+
threshold_ns = int(0.5 * 1_000_000_000)
171+
172+
result = []
173+
buf_words = [] # list of TranscriptWord protobuf
174+
buf_start = None
175+
buf_end = 0
176+
buf_text = []
177+
prev_end = None
178+
179+
for tw in transcript_words:
180+
# Detect word-level time gap
181+
if prev_end is not None and (tw['start'] - prev_end) >= threshold_ns and buf_text:
182+
seg_text = ' '.join(buf_text)
183+
result.append(backend_pb2.TranscriptSegment(
184+
id=len(result),
185+
start=buf_start,
186+
end=buf_end,
187+
text=seg_text,
188+
words=list(buf_words),
189+
))
190+
buf_words = []
191+
buf_text = []
192+
buf_start = None
193+
194+
if buf_start is None:
195+
buf_start = tw['start']
196+
buf_end = tw['end']
197+
buf_text.append(tw['text'])
198+
buf_words.append(backend_pb2.TranscriptWord(
199+
start=tw['start'], end=tw['end'], text=tw['text']
200+
))
201+
prev_end = tw['end']
202+
203+
# flush remaining
204+
if buf_text and buf_start is not None:
205+
seg_text = ' '.join(buf_text)
206+
result.append(backend_pb2.TranscriptSegment(
207+
id=len(result),
208+
start=buf_start,
209+
end=buf_end,
210+
text=seg_text,
211+
words=list(buf_words),
212+
))
213+
214+
return result
215+
87216
def AudioTranscription(self, request, context):
88217
result_segments = []
89218
text = ""
@@ -93,26 +222,67 @@ def AudioTranscription(self, request, context):
93222
print(f"Error: Audio file not found: {audio_path}", file=sys.stderr)
94223
return backend_pb2.TranscriptResult(segments=[], text="")
95224

96-
# NEMO's transcribe method accepts a list of audio paths and returns a list of transcripts
97-
results = self.model.transcribe([audio_path])
225+
# Determine requested timestamp granularity
226+
timestamp_granularities = list(request.timestamp_granularities) if request.timestamp_granularities else []
227+
want_timestamps = bool(timestamp_granularities)
98228

99-
if not results or len(results) == 0:
100-
return backend_pb2.TranscriptResult(segments=[], text="")
229+
if want_timestamps:
230+
# Request timestamps from NeMo.
231+
# timestamps=True forces NeMo to return Hypothesis objects with
232+
# the timestamp dict populated, so we omit return_hypotheses to
233+
# let NeMo choose the correct return type.
234+
results = self.model.transcribe([audio_path], timestamps=True)
101235

102-
# Get the transcript text from the first result.
103-
# CTC models return List[str], TDT/RNNT models return List[Hypothesis]
104-
# where the actual text lives in Hypothesis.text.
105-
result = results[0]
106-
if isinstance(result, str):
107-
text = result
236+
if results and len(results) > 0:
237+
hypotheses = results[0] if isinstance(results[0], list) else results
238+
if hypotheses and len(hypotheses) > 0:
239+
hypothesis = hypotheses[0]
240+
241+
# Hypothesis object should have .timestamp populated
242+
if not hasattr(hypothesis, 'timestamp') or not isinstance(hypothesis.timestamp, dict):
243+
print(
244+
"Warning: timestamps were requested but NeMo did not return "
245+
"Hypothesis objects; falling back to untimestamped output",
246+
file=sys.stderr,
247+
)
248+
249+
# Extract text
250+
if hasattr(hypothesis, 'text'):
251+
text = hypothesis.text or ""
252+
elif isinstance(hypothesis, str):
253+
text = hypothesis
254+
255+
# Build segments with word-level timestamps
256+
stride = self._get_stride_seconds()
257+
result_segments = self._build_segments_with_words(
258+
hypothesis, stride, timestamp_granularities
259+
)
260+
261+
# If no word offsets but we have text, fall back to single segment
262+
if not result_segments and text:
263+
result_segments.append(backend_pb2.TranscriptSegment(
264+
id=0, start=0, end=0, text=text
265+
))
108266
else:
109-
text = getattr(result, 'text', None) or ""
267+
# Simple transcription without timestamps
268+
# NEMO's transcribe method accepts a list of audio paths and returns a list of transcripts
269+
results = self.model.transcribe([audio_path])
110270

111-
if text:
112-
# Create a single segment with the full transcription
113-
result_segments.append(backend_pb2.TranscriptSegment(
114-
id=0, start=0, end=0, text=text
115-
))
271+
if results and len(results) > 0:
272+
# Get the transcript text from the first result.
273+
# CTC models return List[str], TDT/RNNT models return List[Hypothesis]
274+
# where the actual text lives in Hypothesis.text.
275+
result = results[0]
276+
if isinstance(result, str):
277+
text = result
278+
else:
279+
text = getattr(result, 'text', None) or ""
280+
281+
if text:
282+
# Create a single segment with the full transcription
283+
result_segments.append(backend_pb2.TranscriptSegment(
284+
id=0, start=0, end=0, text=text
285+
))
116286

117287
except Exception as err:
118288
print(f"Error in AudioTranscription: {err}", file=sys.stderr)

0 commit comments

Comments
 (0)