adolforeyna/whisper-translation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add safety limits to transcription length and translation tokens to prevent memory exhaustion

Browse Source
This commit is contained in:
Adolfo Reyna
2026-02-26 21:25:57 -05:00
parent d27543022d
commit 37cb45c12f
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
transcribe.py
View File

@@ -103,13 +103,18 @@ def main():
original_text = result['text'].strip() original_text = result['text'].strip()
if original_text: if original_text:
# Limit the input length to avoid memory spikes or model glitches
if len(original_text) > 250:
original_text = original_text[:247] + "..."
print(f"\n[EN]: {original_text}") print(f"\n[EN]: {original_text}")
# 2. Translate to all targets # 2. Translate to all targets
for lang_name, (model, tokenizer) in translation_engines.items(): for lang_name, (model, tokenizer) in translation_engines.items():
inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device) inputs = tokenizer(original_text, return_tensors="pt", padding=True).to(device)
with torch.no_grad(): with torch.no_grad():
translated_tokens = model.generate(**inputs) # Added max_new_tokens to prevent runaway generation
translated_tokens = model.generate(**inputs, max_new_tokens=150)
translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True) translated_text = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
print(f"[{lang_name[:2].upper()}]: {translated_text}") print(f"[{lang_name[:2].upper()}]: {translated_text}")