add the option to use adjacent cues as prompts during transcription

baxtree · Feb 14, 2025 · 6ce24da · 6ce24da
1 parent 5967c8a
commit 6ce24da
Show file tree

Hide file tree

Showing 9 changed files with 66 additions and 88 deletions.
diff --git a/Makefile b/Makefile
@@ -28,28 +28,11 @@ BROWSER := python -c "$$BROWSER_PYSCRIPT"
 
 install:
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
 	.$(PYTHON)/bin/pip install -e . --ignore-installed
-	cp ./bin/subaligner_1pass .$(PYTHON)/bin/subaligner_1pass
-	cp ./bin/subaligner_2pass .$(PYTHON)/bin/subaligner_2pass
-	cp ./bin/subaligner_batch .$(PYTHON)/bin/subaligner_batch
-	cp ./bin/subaligner_convert .$(PYTHON)/bin/subaligner_convert
-	cp ./bin/subaligner_train .$(PYTHON)/bin/subaligner_train
-	cp ./bin/subaligner_tune .$(PYTHON)/bin/subaligner_tune
-	cp ./bin/subaligner .$(PYTHON)/bin/subaligner
+	.$(PYTHON)/bin/pip install -e ".[harmony]"
 
 uninstall:
-	rm -f .$(PYTHON)/bin/subaligner
-	rm -f .$(PYTHON)/bin/subaligner_1pass
-	rm -f .$(PYTHON)/bin/subaligner_2pass
-	rm -f .$(PYTHON)/bin/subaligner_batch
-	rm -f .$(PYTHON)/bin/subaligner_convert
-	rm -f .$(PYTHON)/bin/subaligner_train
-	rm -f .$(PYTHON)/bin/subaligner_tune
+	.$(PYTHON)/bin/pip uninstall subaligner
 
 install-basic:
 	.$(PYTHON)/bin/pip install -e '.' --no-cache-dir
@@ -71,44 +54,33 @@ install-harmony:
 
 build-gzip:
 	mkdir -p dist
-	tar -czf dist/subligner.tar.gz subaligner bin requirements.txt setup.py README.md LICENCE
+	tar -czf dist/subligner.tar.gz subaligner bin pyproject.toml README.md LICENCE
 
 build-rpm:
 	mkdir -p BUILD RPMS SRPMS SOURCES BUILDROOT
-	tar -czf SOURCES/subligner.tar.gz subaligner bin requirements.txt setup.py README.md LICENCE
+	tar -czf SOURCES/subligner.tar.gz subaligner bin pyproject.toml README.md LICENCE
 
 test:
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
 	PYTHONPATH=. .$(PYTHON)/bin/python -m unittest discover
-	-.$(PYTHON)/bin/pycodestyle subaligner tests examples misc bin/subaligner bin/subaligner_1pass bin/subaligner_2pass bin/subaligner_batch bin/subaligner_convert bin/subaligner_train  bin/subaligner_tune setup.py --ignore=E203,E501,W503 --exclude="subaligner/lib"
+	-.$(PYTHON)/bin/pycodestyle subaligner tests examples misc bin/subaligner bin/subaligner_1pass bin/subaligner_2pass bin/subaligner_batch bin/subaligner_convert bin/subaligner_train bin/subaligner_tune setup.py --ignore=E203,E501,E902,W503 --exclude="subaligner/lib"
 
 test-all: ## run tests on every Python version with tox
 	.$(PYTHON)/bin/tox
 
 test-int: ## integration test
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install
 	.$(PYTHON)/bin/pip install -e . --ignore-installed
-	( \
-		source .$(PYTHON)/bin/activate; \
-		radish -b tests/integration/radish tests/integration/feature; \
-	)
+	.$(PYTHON)/bin/pip install -e ".[dev]"
+	source .$(PYTHON)/bin/activate
+	radish -b tests/integration/radish tests/integration/feature;
 
 pydoc: clean-doc ## generate pydoc HTML documentation based on docstrings
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
 	.$(PYTHON)/bin/python -m pydoc -w subaligner; mv subaligner.html docs/index.html
 	.$(PYTHON)/bin/python -m pydoc -w subaligner.embedder; mv subaligner.embedder.html docs
 	.$(PYTHON)/bin/python -m pydoc -w subaligner.exception; mv subaligner.exception.html docs
@@ -128,32 +100,30 @@ pydoc: clean-doc ## generate pydoc HTML documentation based on docstrings
 
 coverage: ## check code coverage quickly with the default Python
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
 	.$(PYTHON)/bin/coverage run --source subaligner -m unittest discover
 	.$(PYTHON)/bin/coverage report
 	.$(PYTHON)/bin/coverage html
 	$(BROWSER) htmlcov/index.html
 
 manual: clean-manual ## generate manual pages
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements-site.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-    SPHINXAPIDOC=../.$(PYTHON)/bin/sphinx-apidoc SPHINXBUILD=../.$(PYTHON)/bin/sphinx-build make -C ./site html; \
-	.$(PYTHON)/bin/python -m sphinx -T -b html -d ./site/build/doctrees -D language=en ./site/source ./site/build/html; \
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
+	SPHINXAPIDOC=../.$(PYTHON)/bin/sphinx-apidoc SPHINXBUILD=../.$(PYTHON)/bin/sphinx-build make -C ./site html
+	.$(PYTHON)/bin/python -m sphinx -T -b html -d ./site/build/doctrees -D language=en ./site/source ./site/build/html
 	$(BROWSER) ./site/build/html/index.html
 
 test-dist:
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
 	.$(PYTHON)/bin/pip install -e .
 
 dist: clean-dist test-dist
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	.$(PYTHON)/bin/python setup.py sdist bdist_wheel
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
+	.$(PYTHON)/bin/pip install build
+	.$(PYTHON)/bin/python -m build --sdist --wheel
 
 release:
 	.$(PYTHON)/bin/twine upload dist/*
@@ -164,22 +134,11 @@ pipenv-install:
 
 profile:
 	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-stretch.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-llm.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
+	.$(PYTHON)/bin/pip install -e . --ignore-installed
+	.$(PYTHON)/bin/pip install -e ".[dev]"
 	.$(PYTHON)/bin/python -c "import misc.profiler; misc.profiler.generate_profiles()"
 	.$(PYTHON)/bin/kernprof -v -l ./misc/profiler.py
 
-app: clean-wheels
-	if [ ! -e ".$(PYTHON)" ]; then ~/.pyenv/versions/$(PYTHON)/bin/python3 -m venv .$(PYTHON); fi
-	.$(PYTHON)/bin/pip install --upgrade pip setuptools wheel; \
-	cat requirements-dev.txt | xargs -L 1 .$(PYTHON)/bin/pip install; \
-	.$(PYTHON)/bin/pip wheel --no-cache-dir --wheel-dir=./wheels -r requirements.txt -r requirements-stretch.txt -r requirements-llm.txt; \
-	STRETCH_OFF=True .$(PYTHON)/bin/python setup.py bdist_wheel -d ./wheels; \
-	.$(PYTHON)/bin/pex subaligner==$(SUBALIGNER_VERSION) --repo=./wheels --platform $(PLATFORM) --no-pypi --no-build --python-shebang="/usr/bin/env python3" -e subaligner -o subaligner-$(PLATFORM).app; \
-
 docker-build:
 	docker build --build-arg RELEASE_VERSION=$(SUBALIGNER_VERSION) -f docker/Dockerfile-Ubuntu20 .
 	docker build --build-arg RELEASE_VERSION=$(SUBALIGNER_VERSION) -f docker/Dockerfile-Ubuntu22 .

diff --git a/README.md b/README.md
@@ -105,18 +105,19 @@ $ subaligner -m transcribe -v video.mp4 -ml eng -mr whisper -mf small -o subtitl
 $ subaligner -m transcribe -v video.mp4 -ml zho -mr whisper -mf medium -o subtitle_aligned.srt
 $ subaligner -m transcribe -v video.mp4 -ml eng -mr whisper -mf turbo -ip "your initial prompt" -o subtitle_aligned.srt
 $ subaligner -m transcribe -v video.mp4 -s subtitle.srt -ml eng -mr whisper -mf turbo -o subtitle_aligned.srt
+$ subaligner -m transcribe -v video.mp4 -s subtitle.srt -upp -ml eng -mr whisper -mf turbo -o subtitle_aligned.srt
 ```
 ```
 # Alignment on segmented plain texts (double newlines as the delimiter)
 
-$ subaligner -m script -v test.mp4 -s subtitle.txt -o subtitle_aligned.srt
+$ subaligner -m script -v video.mp4 -s subtitle.txt -o subtitle_aligned.srt
 $ subaligner -m script -v https://example.com/video.mp4 -s https://example.com/subtitle.txt -o subtitle_aligned.srt
 ```
 ```
 # Alignment on multiple subtitles against the single media file
 
-$ subaligner -m script -v test.mp4 -s subtitle_lang_1.txt -s subtitle_lang_2.txt
-$ subaligner -m script -v test.mp4 -s subtitle_lang_1.txt subtitle_lang_2.txt
+$ subaligner -m script -v video.mp4 -s subtitle_lang_1.txt -s subtitle_lang_2.txt
+$ subaligner -m script -v video.mp4 -s subtitle_lang_1.txt subtitle_lang_2.txt
 ```
 ```
 # Alignment on embedded subtitles
@@ -130,7 +131,7 @@ $ subaligner -m dual -v video.mkv -s embedded:stream_index=0 -o subtitle_aligned
 $ subaligner --languages
 $ subaligner -m single -v video.mp4 -s subtitle.srt -t src,tgt
 $ subaligner -m dual -v video.mp4 -s subtitle.srt -t src,tgt
-$ subaligner -m script -v test.mp4 -s subtitle.txt -o subtitle_aligned.srt -t src,tgt
+$ subaligner -m script -v video.mp4 -s subtitle.txt -o subtitle_aligned.srt -t src,tgt
 $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr helsinki-nlp -o subtitle_aligned.srt -t src,tgt
 $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr facebook-mbart -tf large -o subtitle_aligned.srt -t src,tgt
 $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr facebook-m2m100 -tf small -o subtitle_aligned.srt -t src,tgt

diff --git a/pyproject.toml b/pyproject.toml
@@ -54,6 +54,7 @@ dependencies = [
     "pbr==4.0.2",
     "pkgconfig~=1.5.5",
     "pluggy==0.13.1",
+    "psutil~=7.0.0",
     "pyasn1==0.4.8",
     "pyasn1-modules==0.2.7",
     "pybind11~=2.11.1",

diff --git a/site/source/usage.rst b/site/source/usage.rst
@@ -28,16 +28,17 @@ Make sure you have got the virtual environment activated upfront.
     (.venv) $ subaligner -m transcribe -v video.mp4 -ml zho -mr whisper -mf medium -o subtitle_aligned.srt
     (.venv) $ subaligner -m transcribe -v video.mp4 -ml eng -mr whisper -mf turbo -ip "your initial prompt" -o subtitle_aligned.srt
     (.venv) $ subaligner -m transcribe -v video.mp4 -s subtitle.srt -ml eng -mr whisper -mf turbo -o subtitle_aligned.srt
+    (.venv) $ subaligner -m transcribe -v video.mp4 -s subtitle.srt -upp -ml eng -mr whisper -mf turbo -o subtitle_aligned.srt
 
 **Alignment on segmented plain texts (double newlines as the delimiter)**::
 
-    (.venv) $ subaligner -m script -v test.mp4 -s subtitle.txt -o subtitle_aligned.srt
+    (.venv) $ subaligner -m script -v video.mp4 -s subtitle.txt -o subtitle_aligned.srt
     (.venv) $ subaligner -m script -v https://example.com/video.mp4 -s https://example.com/subtitle.txt -o subtitle_aligned.srt
 
 **Alignment on multiple subtitles against the single media file**::
 
-    (.venv) $ subaligner -m script -v test.mp4 -s subtitle_lang_1.txt -s subtitle_lang_2.txt
-    (.venv) $ subaligner -m script -v test.mp4 -s subtitle_lang_1.txt subtitle_lang_2.txt
+    (.venv) $ subaligner -m script -v video.mp4 -s subtitle_lang_1.txt -s subtitle_lang_2.txt
+    (.venv) $ subaligner -m script -v video.mp4 -s subtitle_lang_1.txt subtitle_lang_2.txt
 
 
 **Alignment on embedded subtitles**::
@@ -50,7 +51,7 @@ Make sure you have got the virtual environment activated upfront.
     (.venv) $ subaligner --languages
     (.venv) $ subaligner -m single -v video.mp4 -s subtitle.srt -t src,tgt
     (.venv) $ subaligner -m dual -v video.mp4 -s subtitle.srt -t src,tgt
-    (.venv) $ subaligner -m script -v test.mp4 -s subtitle.txt -o subtitle_aligned.srt -t src,tgt
+    (.venv) $ subaligner -m script -v video.mp4 -s subtitle.txt -o subtitle_aligned.srt -t src,tgt
     (.venv) $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr helsinki-nlp -o subtitle_aligned.srt -t src,tgt
     (.venv) $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr facebook-mbart -tf large -o subtitle_aligned.srt -t src,tgt
     (.venv) $ subaligner -m dual -v video.mp4 -s subtitle.srt -tr facebook-m2m100 -tf small -o subtitle_aligned.srt -t src,tgt

diff --git a/subaligner/__init__.py b/subaligner/__init__.py
@@ -15,3 +15,5 @@
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 os.environ["TF_CPP_MIN_VLOG_LEVEL"] = "0"
 logging.getLogger("tensorflow").disabled = True
+
+os.environ["TF_USE_LEGACY_KERAS"] = "1"
diff --git a/subaligner/__main__.py b/subaligner/__main__.py
@@ -4,15 +4,15 @@
                   [-sil {afr,amh,ara,arg,asm,aze,ben,bos,bul,cat,ces,cmn,cym,dan,deu,ell,eng,epo,est,eus,fas,fin,fra,gla,gle,glg,grc,grn,guj,heb,hin,hrv,hun,hye,ina,ind,isl,ita,jbo,jpn,kal,kan,kat,kir,kor,kur,lat,lav,lfn,lit,mal,mar,mkd,mlt,msa,mya,nah,nep,nld,nor,ori,orm,pan,pap,pol,por,ron,rus,sin,slk,slv,spa,sqi,srp,swa,swe,tam,tat,tel,tha,tsn,tur,ukr,urd,vie,yue,zho}]
                   [-fos] [-tod TRAINING_OUTPUT_DIRECTORY] [-o OUTPUT] [-t TRANSLATE] [-os OFFSET_SECONDS]
                   [-ml {afr,amh,ara,arg,asm,aze,ben,bos,bul,cat,ces,cmn,cym,dan,deu,ell,eng,epo,est,eus,fas,fin,fra,gla,gle,glg,grc,grn,guj,heb,hin,hrv,hun,hye,ina,ind,isl,ita,jbo,jpn,kal,kan,kat,kir,kor,kur,lat,lav,lfn,lit,mal,mar,mkd,mlt,msa,mya,nah,nep,nld,nor,ori,orm,pan,pap,pol,por,ron,rus,sin,slk,slv,spa,sqi,srp,swa,swe,tam,tat,tel,tha,tsn,tur,ukr,urd,vie,yue,zho}]
-                  [-mr {whisper}] [-mf {tiny,tiny.en,small,medium,medium.en,base,base.en,large-v1,large-v2,large-v3,large,turbo}] [-ip INITIAL_PROMPT] [-tr {helsinki-nlp,whisper,facebook-mbart}]
-                  [-tf TRANSLATION_FLAVOUR] [-mpt MEDIA_PROCESS_TIMEOUT] [-sat SEGMENT_ALIGNMENT_TIMEOUT] [-lgs] [-d] [-q] [-ver]
+                  [-mr {whisper}] [-mf {tiny,tiny.en,small,medium,medium.en,base,base.en,large-v1,large-v2,large-v3,large,turbo}] [-ip INITIAL_PROMPT] [-mcl MAX_CHAR_LENGTH]
+                  [-tr {helsinki-nlp,whisper,facebook-mbart,facebook-m2m100}] [-tf TRANSLATION_FLAVOUR] [-mpt MEDIA_PROCESS_TIMEOUT] [-sat SEGMENT_ALIGNMENT_TIMEOUT] [-wt] [-upp] [-lgs] [-d] [-q] [-ver]
 
 Subaligner command line interface
 
 optional arguments:
   -h, --help            show this help message and exit
   -s SUBTITLE_PATH [SUBTITLE_PATH ...], --subtitle_path SUBTITLE_PATH [SUBTITLE_PATH ...]
-                        File path or URL to the subtitle file (Extensions of supported subtitles: .ass, .smi, .scc, .vtt, .stl, .txt, .sbv, .ssa, .sub, .ttml, .xml, .srt, .ytt, .dfxp, .sami, .tmp) or selector for the embedded subtitle (e.g., embedded:page_num=888 or embedded:stream_index=0)
+                        File path or URL to the subtitle file (Extensions of supported subtitles: .vtt, .srt, .ass, .sbv, .sub, .txt, .ttml, .ssa, .dfxp, .ytt, .stl, .tmp, .smi, .scc, .sami, .xml) or selector for the embedded subtitle (e.g., embedded:page_num=888 or embedded:stream_index=0)
   -l MAX_LOGLOSS, --max_logloss MAX_LOGLOSS
                         Max global log loss for alignment
   -so, --stretch_on     Switch on stretch on subtitles)
@@ -36,14 +36,18 @@
                         Flavour variation for a specific LLM recipe supporting transcription
   -ip INITIAL_PROMPT, --initial_prompt INITIAL_PROMPT
                         Optional text to provide the transcribing context or specific phrases
-  -tr {helsinki-nlp,whisper,facebook-mbart}, --translation_recipe {helsinki-nlp,whisper,facebook-mbart}
+  -mcl MAX_CHAR_LENGTH, --max_char_length MAX_CHAR_LENGTH
+                        Maximum number of characters for each generated subtitle segment
+  -tr {helsinki-nlp,whisper,facebook-mbart,facebook-m2m100}, --translation_recipe {helsinki-nlp,whisper,facebook-mbart,facebook-m2m100}
                         LLM recipe used for translating subtitles
   -tf TRANSLATION_FLAVOUR, --translation_flavour TRANSLATION_FLAVOUR
                         Flavour variation for a specific LLM recipe supporting translation
   -mpt MEDIA_PROCESS_TIMEOUT, --media_process_timeout MEDIA_PROCESS_TIMEOUT
                         Maximum waiting time in seconds when processing media files
   -sat SEGMENT_ALIGNMENT_TIMEOUT, --segment_alignment_timeout SEGMENT_ALIGNMENT_TIMEOUT
                         Maximum waiting time in seconds when aligning each segment
+  -upp, --use_prior_prompting
+                        Whether to use the previous subtitle cue as the current prompt.
   -lgs, --languages     Print out language codes used for stretch and translation
   -d, --debug           Print out debugging information
   -q, --quiet           Switch off logging information
@@ -225,6 +229,8 @@ def main():
         default=60,
         help="Maximum waiting time in seconds when aligning each segment"
     )
+    parser.add_argument("-upp", "--use_prior_prompting", action="store_true",
+                        help="Whether to use the previous subtitle cue as the current prompt.")
     parser.add_argument("-lgs", "--languages", action="store_true",
                         help="Print out language codes used for stretch and translation")
     parser.add_argument("-d", "--debug", action="store_true",
@@ -371,7 +377,8 @@ def main():
                         subtitle, frame_rate = transcriber.transcribe_with_subtitle_as_prompts(video_file_path=local_video_path,
                                                                                                subtitle_file_path=local_subtitle_path,
                                                                                                language_code=stretch_in_lang,
-                                                                                               max_char_length=FLAGS.max_char_length)
+                                                                                               max_char_length=FLAGS.max_char_length,
+                                                                                               use_prior_prompting=FLAGS.use_prior_prompting)
                     aligned_subs = subtitle.subs
                 else:
                     print("ERROR: Unknown mode {}".format(FLAGS.mode))

diff --git a/subaligner/network.py b/subaligner/network.py
@@ -4,10 +4,10 @@
 import psutil
 import numpy as np
 import tensorflow as tf
-import tensorflow.keras.optimizers as tf_optimizers
+import keras.optimizers as tf_optimizers
 
 from typing import Tuple, Optional, List, Generator
-from tensorflow.keras.layers import (
+from keras.layers import (
     Dense,
     Input,
     LSTM,
@@ -19,14 +19,14 @@
     Bidirectional,
 )
 
-from tensorflow.keras.callbacks import (
+from keras.callbacks import (
     EarlyStopping,
     ModelCheckpoint,
     TensorBoard,
     CSVLogger,
 )
-from tensorflow.keras.models import Model, load_model, save_model
-from tensorflow.keras import backend as K
+from keras.models import Model, load_model, save_model
+from keras import backend as K
 from .utils import Utils
 from .logger import Logger
 from .exception import TerminalException