Further refinement of TensorRT-LLM backend based on WhisperS2T

Wordcab · Mar 25, 2024 · 6c81aa3 · 6c81aa3
1 parent dd954a9
commit 6c81aa3
Show file tree

Hide file tree

Showing 10 changed files with 4 additions and 101 deletions.
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/frame_vad_model_cpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/frame_vad_model_cpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/frame_vad_model_gpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/frame_vad_model_gpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/lang_codes.txt b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/lang_codes.txt
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/mel_filters.npz b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/mel_filters.npz
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/seg_vad_model_cpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/seg_vad_model_cpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/seg_vad_model_gpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/seg_vad_model_gpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/silent.mp3 b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/silent.mp3
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/vad_pp_cpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/vad_pp_cpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/vad_pp_gpu.ts b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/assets/vad_pp_gpu.ts
diff --git a/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/create_trt_model.py b/src/wordcab_transcribe/engines/tensorrt_llm/engine_builder/create_trt_model.py
@@ -85,11 +85,11 @@ def build_whisper_trt_model(
     """
     model_url = _MODELS[model_name]
     expected_sha256 = model_url.split("/")[-2]
-    model_ckpt_path = f"../assets/{model_name}.pt"
+    model_ckpt_path = f"assets/{model_name}.pt"
     tokenizer_path = f"{output_dir}/tokenizer.json"
 
     if not os.path.exists(model_ckpt_path):
-        os.makedirs("../assets", exist_ok=True)
+        os.makedirs("assets", exist_ok=True)
 
         logger.info(f"Downloading model '{model_name}' from {model_url}...")
 
@@ -115,7 +115,9 @@ def build_whisper_trt_model(
                     "Model has been downloaded but the SHA256 checksum does not not"
                     " match. Please retry loading the model."
                 )
+        logger.info(f"Model '{model_name}' has been downloaded successfully.")
 
+    print(output_dir, os.path.exists(output_dir))
     if not os.path.exists(output_dir):
         logger.info("Building the model...")
         command = [