Checkpoint current transcription pipeline state

engine_translate.py

@@ -77,23 +77,29 @@ def run_translation(in_queue, out_queue, args):
                 "speaker": item.get("speaker"),
                 "ts": item.get("ts", time.time())
             }
-            if args.en: payload["en"] = corrected_text or raw_text
+            
+            if args.only_translate_llm:
+                text_to_translate = paragraph_text
+                if args.en and paragraph_text:
+                    payload["en"] = paragraph_text
+            else:
+                text_to_translate = paragraph_text or corrected_text or raw_text
+                if args.en:
+                    payload["en"] = corrected_text or raw_text
 
-            text_to_translate = paragraph_text or corrected_text or raw_text
-            if text_to_translate:
-                if translation_engines:
-                    for lang_key, (model, tokenizer) in translation_engines.items():
-                        chunks = split_text_for_translation(text_to_translate, args.mt_max_chars)
-                        translated_parts = []
-                        for chunk in chunks:
-                            inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
-                            with torch.no_grad():
-                                translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
-                            translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
-                        payload[lang_key] = " ".join(translated_parts)
-                
-                # ALWAYS put in queue, even if no translations were done
-                out_queue.put(payload)
+            if text_to_translate and translation_engines:
+                for lang_key, (model, tokenizer) in translation_engines.items():
+                    chunks = split_text_for_translation(text_to_translate, args.mt_max_chars)
+                    translated_parts = []
+                    for chunk in chunks:
+                        inputs = tokenizer(chunk, return_tensors="pt", padding=True).to(device)
+                        with torch.no_grad():
+                            translated_tokens = model.generate(**inputs, max_new_tokens=args.mt_max_new_tokens, num_beams=args.mt_num_beams)
+                        translated_parts.append(tokenizer.decode(translated_tokens[0], skip_special_tokens=True).strip())
+                    payload[lang_key] = " ".join(translated_parts)
+            
+            # ALWAYS put in queue, even if no translations were done
+            out_queue.put(payload)
         except Exception as e: print(f"[Translate] Error: {e}")
 
 if __name__ == "__main__":