From f7b8e7b1bd0864c7d5c7e7cf62d1dcc40b3eee1b Mon Sep 17 00:00:00 2001
From: philmcmahon <philip.mcmahon@theguardian.com>
Date: Tue, 21 Jan 2025 21:00:47 +0000
Subject: [PATCH] Add script to download whisperx models

---
 roles/whisperx/tasks/main.yml | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/roles/whisperx/tasks/main.yml b/roles/whisperx/tasks/main.yml
index e9f07836..e619b6fd 100644
--- a/roles/whisperx/tasks/main.yml
+++ b/roles/whisperx/tasks/main.yml
@@ -19,3 +19,26 @@
     name: whisperx
     executable: pip3
 
+- name: Install dependencies needed for downloading models
+  pip:
+    name:
+     - huggingface_hub
+     - typer
+    executable: pip3
+
+- name: Download model fetch script
+  get_url:
+      url: "https://raw.githubusercontent.com/guardian/transcription-service/refs/heads/add-whisperx-support/scripts/download_whisperx_models.py"
+      dest: "/tmp/download_whisperx_models.py"
+      mode: "0755"
+
+- name: Download whisperx models
+  command: "python3 /tmp/download_whisperx_models.py --whisper-models --diarization-models --torch-align-models --huggingface-token {{ huggingface_token }}"
+
+- name: Remove download script dependencies
+  pip:
+    name:
+     - huggingface_hub
+     - typer
+    state: absent
+    executable: pip3