mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
707 lines
25 KiB
Python
707 lines
25 KiB
Python
"""FunASR CLI - Agent-friendly speech recognition from the command line."""
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import re
|
||
import sys
|
||
import time
|
||
import unicodedata
|
||
|
||
MODEL_CONFIGS = {
|
||
"sensevoice": {"model": "iic/SenseVoiceSmall", "vad_model": "fsmn-vad", "vad_kwargs": {"max_single_segment_time": 30000}},
|
||
"paraformer": {"model": "paraformer-zh", "vad_model": "fsmn-vad", "punc_model": "ct-punc"},
|
||
"paraformer-en": {"model": "paraformer-en", "vad_model": "fsmn-vad"},
|
||
"fun-asr-nano": {"model": "FunAudioLLM/Fun-ASR-Nano-2512", "vad_model": "fsmn-vad"},
|
||
}
|
||
|
||
SUBTITLE_CONTINUATION_PUNCTUATION = (
|
||
",",
|
||
",",
|
||
"、",
|
||
":",
|
||
":",
|
||
";",
|
||
";",
|
||
)
|
||
SUBTITLE_TRAILING_PUNCTUATION = SUBTITLE_CONTINUATION_PUNCTUATION + (
|
||
"。",
|
||
".",
|
||
"!",
|
||
"!",
|
||
"?",
|
||
"?",
|
||
)
|
||
|
||
|
||
def clean_text(text):
|
||
return re.sub(r"<\|[^|]*\|>", "", text).strip()
|
||
|
||
|
||
def _srt_time(ms):
|
||
ms = max(0, int(round(ms)))
|
||
h = ms // 3600000
|
||
m = (ms % 3600000) // 60000
|
||
sec = (ms % 60000) // 1000
|
||
ms_rem = ms % 1000
|
||
return f"{h:02d}:{m:02d}:{sec:02d},{ms_rem:03d}"
|
||
|
||
|
||
def format_srt(segments):
|
||
lines = []
|
||
for i, seg in enumerate(segments, 1):
|
||
lines += [str(i), f"{_srt_time(seg.get('start',0))} --> {_srt_time(seg.get('end',0))}", seg.get('text',''), ""]
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _subtitle_body_length(text):
|
||
return len(
|
||
str(text).strip().rstrip("".join(SUBTITLE_TRAILING_PUNCTUATION)).strip()
|
||
)
|
||
|
||
|
||
def _join_subtitle_text(left, right):
|
||
left = str(left).rstrip()
|
||
right = str(right).lstrip()
|
||
if (
|
||
left
|
||
and right
|
||
and left[-1].isascii()
|
||
and right[0].isascii()
|
||
and left[-1].isalnum()
|
||
and right[0].isalnum()
|
||
):
|
||
return f"{left} {right}"
|
||
return left + right
|
||
|
||
|
||
def _subtitle_token_spans(text):
|
||
spans = []
|
||
pending_start = None
|
||
index = 0
|
||
while index < len(text):
|
||
char = text[index]
|
||
if char.isspace():
|
||
index += 1
|
||
continue
|
||
if unicodedata.category(char).startswith("P"):
|
||
if spans:
|
||
spans[-1][1] = index + 1
|
||
elif pending_start is None:
|
||
pending_start = index
|
||
index += 1
|
||
continue
|
||
|
||
start = index
|
||
if char.isascii() and (char.isalnum() or char in "_'"):
|
||
index += 1
|
||
while index < len(text):
|
||
char = text[index]
|
||
if not (char.isascii() and (char.isalnum() or char in "_'")):
|
||
break
|
||
index += 1
|
||
elif _is_supported_subtitle_character(char):
|
||
index += 1
|
||
else:
|
||
return []
|
||
if pending_start is not None:
|
||
start = pending_start
|
||
pending_start = None
|
||
spans.append([start, index])
|
||
|
||
if pending_start is not None and spans:
|
||
spans[-1][1] = len(text)
|
||
return spans
|
||
|
||
|
||
def _is_supported_subtitle_character(char):
|
||
codepoint = ord(char)
|
||
return (
|
||
0x3400 <= codepoint <= 0x4DBF
|
||
or 0x4E00 <= codepoint <= 0x9FFF
|
||
or 0xF900 <= codepoint <= 0xFAFF
|
||
or 0x20000 <= codepoint <= 0x323AF
|
||
or 0x3040 <= codepoint <= 0x30FF
|
||
or 0x31F0 <= codepoint <= 0x31FF
|
||
or 0xFF66 <= codepoint <= 0xFF9D
|
||
or 0x1100 <= codepoint <= 0x11FF
|
||
or 0x3130 <= codepoint <= 0x318F
|
||
or 0xAC00 <= codepoint <= 0xD7AF
|
||
)
|
||
|
||
|
||
def _subtitle_word_spans(text, words):
|
||
spans = []
|
||
cursor = 0
|
||
for raw_word in words:
|
||
word = str(raw_word).lstrip("▁").strip()
|
||
if not word:
|
||
return []
|
||
start = text.find(word, cursor)
|
||
if start < 0 or any(
|
||
not (char.isspace() or unicodedata.category(char).startswith("P"))
|
||
for char in text[cursor:start]
|
||
):
|
||
return []
|
||
if spans:
|
||
spans[-1][1] = start
|
||
elif any(
|
||
not (char.isspace() or unicodedata.category(char).startswith("P"))
|
||
for char in text[:start]
|
||
):
|
||
return []
|
||
end = start + len(word)
|
||
spans.append([0 if not spans and start else start, end])
|
||
cursor = end
|
||
|
||
if any(
|
||
not (char.isspace() or unicodedata.category(char).startswith("P"))
|
||
for char in text[cursor:]
|
||
):
|
||
return []
|
||
if spans:
|
||
spans[-1][1] = len(text)
|
||
return spans
|
||
|
||
|
||
def _timestamp_pair(item):
|
||
if not isinstance(item, (list, tuple)) or len(item) < 2:
|
||
return None
|
||
try:
|
||
start = int(item[0])
|
||
end = int(item[1])
|
||
except (TypeError, ValueError, OverflowError):
|
||
return None
|
||
return [start, end] if end > start else None
|
||
|
||
|
||
def _timestamps_are_ordered(timestamps):
|
||
return bool(timestamps) and all(
|
||
timestamp is not None
|
||
and timestamp[0] >= 0
|
||
and (index == 0 or timestamp[0] >= timestamps[index - 1][1])
|
||
for index, timestamp in enumerate(timestamps)
|
||
)
|
||
|
||
|
||
def _subtitle_break_weights(text, token_spans):
|
||
"""Map token-boundary indices to lexical and punctuation preferences."""
|
||
boundary_to_token = {span[1]: index + 1 for index, span in enumerate(token_spans)}
|
||
breaks = {len(token_spans): 0.0}
|
||
|
||
try:
|
||
import logging
|
||
import warnings
|
||
|
||
with warnings.catch_warnings():
|
||
warnings.filterwarnings(
|
||
"ignore",
|
||
message="pkg_resources is deprecated as an API.*",
|
||
category=UserWarning,
|
||
)
|
||
import jieba
|
||
|
||
jieba.setLogLevel(logging.ERROR)
|
||
cursor = 0
|
||
for piece in jieba.cut(text, cut_all=False, HMM=True):
|
||
cursor += len(piece)
|
||
if piece.strip() and any(
|
||
char.isalnum() or _is_supported_subtitle_character(char)
|
||
for char in piece
|
||
):
|
||
token_index = boundary_to_token.get(cursor)
|
||
if token_index is not None:
|
||
body_length = sum(
|
||
not (
|
||
char.isspace()
|
||
or unicodedata.category(char).startswith("P")
|
||
)
|
||
for char in piece
|
||
)
|
||
strength = 1.0 if body_length > 1 else 0.2
|
||
breaks[token_index] = max(
|
||
breaks.get(token_index, 0.0), strength
|
||
)
|
||
except (ImportError, RuntimeError, ValueError):
|
||
pass
|
||
|
||
for index, span in enumerate(token_spans[:-1], 1):
|
||
boundary = span[1]
|
||
left = text[:boundary].rstrip()
|
||
right = text[boundary:].lstrip()
|
||
strength = 0.0
|
||
if left and left[-1] in ".!?。!?;;":
|
||
strength = 4.0
|
||
elif left and unicodedata.category(left[-1]).startswith("P"):
|
||
strength = 2.0
|
||
elif (
|
||
text[boundary - 1 : boundary].isspace()
|
||
or text[boundary : boundary + 1].isspace()
|
||
):
|
||
strength = 1.5
|
||
elif left and right and left[-1].isascii() != right[0].isascii():
|
||
strength = 1.0
|
||
if strength:
|
||
breaks[index] = max(breaks.get(index, 0.0), strength)
|
||
return breaks
|
||
|
||
|
||
def _balanced_subtitle_ranges(
|
||
text, token_spans, timestamps, max_duration_ms, max_chars
|
||
):
|
||
"""Find the fewest valid cues, then optimize their readable boundaries."""
|
||
token_count = len(token_spans)
|
||
if not token_count:
|
||
return []
|
||
|
||
total_chars = len(text.strip())
|
||
total_duration = sum(timestamp[1] - timestamp[0] for timestamp in timestamps)
|
||
min_cues = max(1, (total_chars + max_chars - 1) // max_chars)
|
||
break_weights = _subtitle_break_weights(text, token_spans)
|
||
|
||
# Dynamic programming avoids the short final fragment produced by a greedy split.
|
||
for cue_count in range(min_cues, token_count + 1):
|
||
target_chars = total_chars / cue_count
|
||
target_duration = total_duration / cue_count
|
||
states = {0: (0.0, [])}
|
||
for cue_index in range(cue_count):
|
||
next_states = {}
|
||
cues_left = cue_count - cue_index - 1
|
||
for token_start, (cost, path) in states.items():
|
||
min_end = token_start + 1
|
||
max_end = token_count - cues_left
|
||
for token_end in range(min_end, max_end + 1):
|
||
cue_text = text[
|
||
token_spans[token_start][0] : token_spans[token_end - 1][1]
|
||
].strip()
|
||
cue_duration = timestamps[token_end - 1][1] - timestamps[token_start][0]
|
||
if len(cue_text) > max_chars or cue_duration > max_duration_ms:
|
||
if token_end > min_end:
|
||
break
|
||
continue
|
||
|
||
break_strength = break_weights.get(token_end)
|
||
unsafe_break = token_end < token_count and break_strength is None
|
||
char_error = (len(cue_text) - target_chars) / max(target_chars, 1.0)
|
||
duration_error = (cue_duration - target_duration) / max(
|
||
target_duration, 1.0
|
||
)
|
||
gap_ms = (
|
||
timestamps[token_end][0] - timestamps[token_end - 1][1]
|
||
if token_end < token_count
|
||
else 0
|
||
)
|
||
cue_cost = (
|
||
(1000.0 if unsafe_break else 0.0)
|
||
+ 8.0 * char_error * char_error
|
||
+ 2.0 * duration_error * duration_error
|
||
- 2.0 * (break_strength or 0.0)
|
||
- min(max(gap_ms, 0), 1000) / 1000.0
|
||
)
|
||
candidate = (cost + cue_cost, [*path, (token_start, token_end)])
|
||
previous = next_states.get(token_end)
|
||
if previous is None or candidate[0] < previous[0]:
|
||
next_states[token_end] = candidate
|
||
states = next_states
|
||
if not states:
|
||
break
|
||
|
||
if token_count in states:
|
||
return states[token_count][1]
|
||
return []
|
||
|
||
|
||
def _sentence_timestamp_words(result):
|
||
sentence_info = result.get("sentence_info", []) or []
|
||
words = result.get("words", []) or []
|
||
raw_timestamps = result.get("timestamp") or result.get("timestamps") or []
|
||
timestamps = [_timestamp_pair(item) for item in raw_timestamps]
|
||
if not words or len(words) != len(timestamps) or not _timestamps_are_ordered(
|
||
timestamps
|
||
):
|
||
return [None] * len(sentence_info)
|
||
|
||
mapped_words = []
|
||
cursor = 0
|
||
for sentence in sentence_info:
|
||
local_timestamps = [
|
||
_timestamp_pair(item)
|
||
for item in (
|
||
sentence.get("timestamp") or sentence.get("timestamps") or []
|
||
)
|
||
]
|
||
if not _timestamps_are_ordered(local_timestamps):
|
||
mapped_words.append(None)
|
||
continue
|
||
|
||
local_cursor = cursor
|
||
selected = []
|
||
for timestamp in local_timestamps:
|
||
while (
|
||
local_cursor < len(timestamps)
|
||
and timestamps[local_cursor] != timestamp
|
||
):
|
||
local_cursor += 1
|
||
if local_cursor == len(timestamps):
|
||
selected = []
|
||
break
|
||
selected.append(words[local_cursor])
|
||
local_cursor += 1
|
||
if len(selected) == len(local_timestamps):
|
||
mapped_words.append(selected)
|
||
cursor = local_cursor
|
||
else:
|
||
mapped_words.append(None)
|
||
return mapped_words
|
||
|
||
|
||
def _split_subtitle_segment(segment, max_duration_ms, max_chars):
|
||
text = str(segment.get("text", ""))
|
||
start = int(segment.get("start", 0) or 0)
|
||
end = int(segment.get("end", start) or start)
|
||
if not text or (end - start <= max_duration_ms and len(text) <= max_chars):
|
||
return [dict(segment)]
|
||
|
||
raw_timestamps = segment.get("timestamp") or segment.get("timestamps") or []
|
||
timestamps = [_timestamp_pair(item) for item in raw_timestamps]
|
||
if not _timestamps_are_ordered(timestamps):
|
||
return [dict(segment)]
|
||
|
||
words = segment.get("words") or []
|
||
token_spans = (
|
||
_subtitle_word_spans(text, words) if words else _subtitle_token_spans(text)
|
||
)
|
||
if not timestamps or len(timestamps) != len(token_spans):
|
||
return [dict(segment)]
|
||
for index, span in enumerate(token_spans):
|
||
token_text = text[span[0] : span[1]].strip()
|
||
if (
|
||
timestamps[index][1] - timestamps[index][0] > max_duration_ms
|
||
or len(token_text) > max_chars
|
||
):
|
||
return [dict(segment)]
|
||
|
||
ranges = _balanced_subtitle_ranges(
|
||
text, token_spans, timestamps, max_duration_ms, max_chars
|
||
)
|
||
if not ranges:
|
||
return [dict(segment)]
|
||
|
||
cues = []
|
||
for token_start, token_end in ranges:
|
||
cue = dict(segment)
|
||
cue["text"] = text[
|
||
token_spans[token_start][0] : token_spans[token_end - 1][1]
|
||
].strip()
|
||
cue["start"] = timestamps[token_start][0]
|
||
cue["end"] = timestamps[token_end - 1][1]
|
||
cue["timestamp"] = timestamps[token_start:token_end]
|
||
cue.pop("timestamps", None)
|
||
cue.pop("words", None)
|
||
cues.append(cue)
|
||
|
||
return cues
|
||
|
||
|
||
def merge_subtitle_segments(
|
||
segments, max_gap_ms=500, max_duration_ms=8000, max_chars=42
|
||
):
|
||
"""Group sentence timestamps into bounded, readable subtitle cues."""
|
||
def can_follow(left, right):
|
||
left_text = str(left.get("text", ""))
|
||
right_text = str(right.get("text", ""))
|
||
gap_ms = right.get("start", 0) - left.get("end", 0)
|
||
left_speaker = left.get("speaker", left.get("spk"))
|
||
right_speaker = right.get("speaker", right.get("spk"))
|
||
return (
|
||
left_speaker == right_speaker
|
||
and 0 <= gap_ms <= max_gap_ms
|
||
and (
|
||
left_text.rstrip().endswith(SUBTITLE_CONTINUATION_PUNCTUATION)
|
||
or _subtitle_body_length(left_text) <= 2
|
||
or (
|
||
gap_ms <= min(max_gap_ms, 100)
|
||
and _subtitle_body_length(right_text) > 2
|
||
and right_text.rstrip().endswith(
|
||
SUBTITLE_CONTINUATION_PUNCTUATION
|
||
)
|
||
)
|
||
)
|
||
)
|
||
|
||
def combine(group):
|
||
cue = dict(group[0])
|
||
cue["end"] = group[-1].get("end", cue.get("end", 0))
|
||
text = str(group[0].get("text", ""))
|
||
for item in group[1:]:
|
||
text = _join_subtitle_text(text, item.get("text", ""))
|
||
cue["text"] = text
|
||
if any(item.get("timestamp") for item in group):
|
||
cue["timestamp"] = [
|
||
timestamp
|
||
for item in group
|
||
for timestamp in item.get("timestamp", [])
|
||
]
|
||
if len(group) > 1 and any(item.get("words") for item in group):
|
||
if all(
|
||
isinstance(item.get("words"), list)
|
||
and item["words"]
|
||
and len(item["words"]) == len(item.get("timestamp", []))
|
||
for item in group
|
||
):
|
||
cue["words"] = [word for item in group for word in item["words"]]
|
||
else:
|
||
cue.pop("words", None)
|
||
return cue
|
||
|
||
def pack(chain):
|
||
groups = []
|
||
current = [chain[-1]]
|
||
for item in reversed(chain[:-1]):
|
||
candidate = [item, *current]
|
||
combined = combine(candidate)
|
||
duration_ms = combined.get("end", 0) - combined.get("start", 0)
|
||
if (
|
||
duration_ms <= max_duration_ms
|
||
and len(combined.get("text", "")) <= max_chars
|
||
):
|
||
current = candidate
|
||
else:
|
||
groups.append(current)
|
||
current = [item]
|
||
groups.append(current)
|
||
return [combine(group) for group in reversed(groups)]
|
||
|
||
merged = []
|
||
chain = []
|
||
for source in segments:
|
||
for current in _split_subtitle_segment(
|
||
source, max_duration_ms=max_duration_ms, max_chars=max_chars
|
||
):
|
||
if chain and not can_follow(chain[-1], current):
|
||
merged.extend(pack(chain))
|
||
chain = []
|
||
chain.append(current)
|
||
if chain:
|
||
merged.extend(pack(chain))
|
||
return merged
|
||
|
||
|
||
def format_tsv(segments):
|
||
lines = ["start\tend\ttext"]
|
||
for seg in segments:
|
||
lines.append(f"{seg.get('start',0)/1000:.3f}\t{seg.get('end',0)/1000:.3f}\t{seg.get('text','')}")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _parse_ms(value, scale=1):
|
||
if value is None:
|
||
return None
|
||
try:
|
||
return int(float(value) * scale)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def _timestamp_bounds_ms(result):
|
||
bounds = []
|
||
for key in ("timestamp", "timestamps"):
|
||
for ts in result.get(key, []) or []:
|
||
if isinstance(ts, dict):
|
||
start = ts.get("start_time", ts.get("start"))
|
||
end = ts.get("end_time", ts.get("end"))
|
||
start_ms = _parse_ms(start, 1000)
|
||
end_ms = _parse_ms(end, 1000)
|
||
elif isinstance(ts, (list, tuple)) and len(ts) >= 2:
|
||
start_ms = _parse_ms(ts[0])
|
||
end_ms = _parse_ms(ts[1])
|
||
else:
|
||
continue
|
||
if start_ms is None or end_ms is None:
|
||
continue
|
||
if end_ms > start_ms:
|
||
bounds.append((start_ms, end_ms))
|
||
if not bounds:
|
||
return None
|
||
return min(start for start, _ in bounds), max(end for _, end in bounds)
|
||
|
||
|
||
def _format_output(text, segments, timestamps, fmt, audio_path, model_name, language, elapsed):
|
||
if fmt == "text":
|
||
return text
|
||
elif fmt == "json":
|
||
obj = {"text": text}
|
||
if segments:
|
||
obj["segments"] = segments
|
||
if timestamps:
|
||
obj["timestamps"] = timestamps
|
||
try:
|
||
import soundfile as sf
|
||
audio_dur = round(sf.info(audio_path).duration, 3)
|
||
except Exception:
|
||
audio_dur = None
|
||
obj.update({"file": os.path.basename(audio_path), "model": model_name, "language": language or "auto", "audio_duration_s": audio_dur, "processing_s": round(elapsed, 3)})
|
||
return json.dumps(obj, ensure_ascii=False, indent=2)
|
||
elif fmt == "srt":
|
||
if segments:
|
||
return format_srt(segments)
|
||
# No per-sentence timestamps: emit one valid cue spanning the known
|
||
# timestamp/audio bounds instead of a bogus 99:59:59 end time.
|
||
timestamp_bounds = _timestamp_bounds_ms({"timestamp": timestamps})
|
||
if timestamp_bounds:
|
||
start_ms, end_ms = timestamp_bounds
|
||
return f"1\n{_srt_time(start_ms)} --> {_srt_time(end_ms)}\n{text}\n"
|
||
try:
|
||
import soundfile as sf
|
||
dur_ms = int(sf.info(audio_path).duration * 1000)
|
||
except Exception:
|
||
dur_ms = 0
|
||
return f"1\n00:00:00,000 --> {_srt_time(dur_ms)}\n{text}\n"
|
||
elif fmt == "tsv":
|
||
return format_tsv(segments) if segments else f"start\tend\ttext\n0.000\t0.000\t{text}"
|
||
|
||
|
||
def _get_version():
|
||
try:
|
||
from funasr import __version__
|
||
return __version__
|
||
except Exception:
|
||
return "unknown"
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(
|
||
prog="funasr",
|
||
description="FunASR - speech recognition CLI. 50+ languages, speaker diarization.",
|
||
epilog="Examples:\n"
|
||
" funasr audio.wav\n"
|
||
" funasr audio.wav --model sensevoice -f json\n"
|
||
" funasr audio.wav -f srt -o ./subs\n"
|
||
" funasr audio.wav --spk --timestamps\n"
|
||
" funasr audio.wav --hub hf --model fun-asr-nano\n",
|
||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
)
|
||
p.add_argument("audio", nargs="+", help="Audio file(s) to transcribe")
|
||
p.add_argument("--model", "-m", default="sensevoice", choices=list(MODEL_CONFIGS), help="Model (default: sensevoice)")
|
||
p.add_argument("--hub", "-H", default="ms", choices=["ms", "hf"], help="Model hub: ms (ModelScope) or hf (Hugging Face). Default: ms")
|
||
p.add_argument("--language", "-l", default=None, help="Language: zh, en, ja, ko, yue, auto")
|
||
p.add_argument("--device", default=None, help="Device: cuda:0, cpu (default: auto)")
|
||
p.add_argument("--output-format", "-f", default="text", choices=["text", "json", "srt", "tsv"], help="Output format (default: text)")
|
||
p.add_argument(
|
||
"--subtitle-segment-mode",
|
||
choices=["readable", "sentence"],
|
||
default="readable",
|
||
help="SRT cue grouping: readable (default) or raw model sentence boundaries",
|
||
)
|
||
p.add_argument("--output-dir", "-o", default=None, help="Write output files to directory")
|
||
p.add_argument("--timestamps", action="store_true", help="Include word-level timestamps")
|
||
p.add_argument("--spk", action="store_true", help="Enable speaker diarization")
|
||
p.add_argument("--hotwords", default=None, help="Comma-separated hotwords")
|
||
p.add_argument("--verbose", "-v", action="store_true", help="Show loading/timing info on stderr")
|
||
p.add_argument("--version", action="version", version=f"%(prog)s {_get_version()}")
|
||
args = p.parse_args()
|
||
|
||
if args.verbose:
|
||
print(f"Loading model: {args.model} ...", file=sys.stderr)
|
||
|
||
import torch
|
||
from funasr import AutoModel
|
||
|
||
device = args.device or ("cuda:0" if torch.cuda.is_available() else "cpu")
|
||
config = MODEL_CONFIGS[args.model].copy()
|
||
config["hub"] = args.hub
|
||
if args.spk and "spk_model" not in config:
|
||
config["spk_model"] = "cam++"
|
||
if "punc_model" not in config and args.model != "fun-asr-nano":
|
||
if args.model != "sensevoice" or args.output_format in ("srt", "tsv"):
|
||
config["punc_model"] = "ct-punc"
|
||
|
||
t_load = time.time()
|
||
model = AutoModel(device=device, disable_update=True, **config)
|
||
if args.verbose:
|
||
print(f"Model loaded in {time.time() - t_load:.1f}s", file=sys.stderr)
|
||
|
||
if args.output_dir:
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
|
||
for audio_path in args.audio:
|
||
if not os.path.isfile(audio_path):
|
||
print(f"Error: file not found: {audio_path}", file=sys.stderr)
|
||
sys.exit(1)
|
||
|
||
if args.verbose:
|
||
print(f"Transcribing: {audio_path}", file=sys.stderr)
|
||
|
||
t0 = time.time()
|
||
gen_kw = {"input": audio_path, "batch_size": 1}
|
||
if args.language:
|
||
gen_kw["language"] = args.language
|
||
if args.hotwords:
|
||
hotwords = [
|
||
word.strip() for word in args.hotwords.split(",") if word.strip()
|
||
]
|
||
if args.model == "paraformer":
|
||
gen_kw["hotword"] = " ".join(hotwords)
|
||
else:
|
||
gen_kw["hotwords"] = hotwords
|
||
|
||
if args.output_format in ("srt", "tsv"):
|
||
gen_kw.update(
|
||
{
|
||
"sentence_timestamp": True,
|
||
"output_timestamp": True,
|
||
"return_time_stamps": True,
|
||
}
|
||
)
|
||
|
||
result = model.generate(**gen_kw)
|
||
elapsed = time.time() - t0
|
||
|
||
text = clean_text(result[0].get("text", ""))
|
||
segments = []
|
||
if "sentence_info" in result[0]:
|
||
sentence_words = _sentence_timestamp_words(result[0])
|
||
for index, seg in enumerate(result[0]["sentence_info"]):
|
||
s = {
|
||
"start": seg.get("start", 0),
|
||
"end": seg.get("end", 0),
|
||
"text": clean_text(seg.get("sentence") or seg.get("text", "")),
|
||
"timestamp": seg.get("timestamp") or seg.get("timestamps"),
|
||
}
|
||
if (
|
||
args.output_format == "srt"
|
||
and args.subtitle_segment_mode == "readable"
|
||
and sentence_words[index]
|
||
):
|
||
s["words"] = sentence_words[index]
|
||
if args.spk and "spk" in seg:
|
||
s["speaker"] = seg["spk"]
|
||
segments.append(s)
|
||
if (
|
||
args.output_format == "srt"
|
||
and args.subtitle_segment_mode == "readable"
|
||
):
|
||
segments = merge_subtitle_segments(segments)
|
||
|
||
timestamps = result[0].get("timestamps") or result[0].get("timestamp")
|
||
if not args.timestamps and args.output_format not in ("srt", "tsv"):
|
||
timestamps = None
|
||
output = _format_output(text, segments, timestamps, args.output_format, audio_path, args.model, args.language, elapsed)
|
||
|
||
if args.output_dir:
|
||
ext = {"text": "txt", "json": "json", "srt": "srt", "tsv": "tsv"}[args.output_format]
|
||
out_path = os.path.join(args.output_dir, os.path.splitext(os.path.basename(audio_path))[0] + "." + ext)
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
f.write(output)
|
||
if args.verbose:
|
||
print(f"Written: {out_path}", file=sys.stderr)
|
||
else:
|
||
print(output)
|
||
|
||
if args.verbose:
|
||
print(f"Done in {elapsed:.2f}s", file=sys.stderr)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|