From d6da26607ede62d8363d769adab1090a831a9dd1 Mon Sep 17 00:00:00 2001
From: Izaya Orihara <mc.duff@outlook.de>
Date: Mon, 9 Oct 2023 12:12:02 +0200
Subject: [PATCH] Feature request: added huggingface functionality checking for
 token

---
 .gitignore         |  5 ++++-
 README.md          | 56 ++++++++++++++++++++++++++++++++++++++++++++--
 src/main.py        |  7 ++++++
 src/transcriber.py |  4 +++-
 4 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 200b914..e95ff19 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,4 +10,7 @@ __pycache__/
 .venv
 env/
 venv/
-ENV/
\ No newline at end of file
+ENV/
+
+# Large files
+data/models # local storage of transcriber models
\ No newline at end of file
diff --git a/README.md b/README.md
index a2f259b..f48317b 100644
--- a/README.md
+++ b/README.md
@@ -4,10 +4,61 @@
 
 Transcriber code
 
+## Setup
+
+Before running this script, make sure to install the necessary python packages using the following command:
+
+```bash
+pip install -r requirements.txt
+```
+
+## Running the script
+
+To generate transcripts create a new job with your specifications (see jobs section below).
+<br/>
+If you need to download the transcriber model (not saved in `data/models`yet) set your huggingface token:
+
+On Linux/Mac:
+
+```bash
+export HF_TOKEN="your-token-here"
+```
+
+On Windows (Command prompt):
+
+```cmd
+set HF_TOKEN="your-token-here"
+```
+
+Then run the following command to generate your transcript:
+
+```bash
+main.py --job {name}
+```
+
+where name is the name of your job.
+
+## Jobs
+
+Jobs specify tasks for transcript creation. A job contains the following information.
+
+- `name`: string: The name of the job
+- `source`: List[string]: The sources of the videos (paths, YouTube links, TibAV links)
+- `modelpath`: Optional[string]: The path to the model if already saved locally (will be downloaded temporarily if not)
+- `language`: Optional[string]: The language of the video (If not specified employs auto-detection)
+
 ## to-do
 
 - seperate data production form code
-- seperate huggingface (hf) token out
+- remember to accept the conditions on diarization if you use huggingface
+- do a config json
+- do option to have models already stored and point to them (is this necessary?)
+    - also for single language models
+- check warnings
+- downloader for tib-av
+- more optional setting
+
+## Structure
 
 ```
 project_root/
@@ -23,6 +74,7 @@ project_root/
 │   └── test_transcript_generator.py
 ├── data/
 │   ├── audio/   # downloaded mp4 and output wav files
+│   ├── models/   # models used for transcripts (are ignored by git and not necessarily)
 │   └── transcripts/  # JSON and text transcript files
 └── README.md
 ```
@@ -37,4 +89,4 @@ project_root/
 - **tests/**: This directory contains all your test scripts corresponding to your source files.
 
 - **data/**: This directory contains all the data needed for your code. For simplicity, I categorize the files into
-  separate directories for audio and transcript files.
\ No newline at end of file
+  separate directories for audio, model and transcript files.
\ No newline at end of file
diff --git a/src/main.py b/src/main.py
index 50e0033..60023fd 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,3 +1,5 @@
+import os
+
 from src.transcript_generator import TranscriptGenerator
 
 
@@ -7,6 +9,11 @@ def main():
         'https://www.youtube.com/watch?v=BVVo1BaL6JA'
     ]
 
+    hf_token = os.environ.get('HF_TOKEN')
+    if not hf_token:
+        hf_token = input("Please enter your HF_TOKEN: ")
+    os.environ['HF_TOKEN'] = hf_token
+
     transcript_generator = TranscriptGenerator()
     transcript_generator.generate(input_sources)
 
diff --git a/src/transcriber.py b/src/transcriber.py
index d73592b..742b4c7 100644
--- a/src/transcriber.py
+++ b/src/transcriber.py
@@ -1,6 +1,7 @@
 import gc
 import json
 import logging
+import os
 
 import torch
 import whisperx
@@ -30,6 +31,7 @@ class Transcriber:
             logger.exception("Failed to initialize models: %s", e)
             raise
         self.transcript_output = None
+        self.hf_token = os.environ.get('HF_TOKEN')
 
     def transcribe(self, wav_filename: str, batch_size: int = 16) -> dict:
         """
@@ -53,7 +55,7 @@ class Transcriber:
             self._clear_memory()
 
             # Run Diarization
-            diarize_model = whisperx.DiarizationPipeline(use_auth_token='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',
+            diarize_model = whisperx.DiarizationPipeline(use_auth_token=self.hf_token,
                                                          device=self.device)
             diarize_segments = diarize_model(audio)
             result = whisperx.assign_word_speakers(diarize_segments, result)