Write current package version to metadata

1 files changed, 5 insertions, 0 deletions

diff --git a/bicleaner_ai/training.py b/bicleaner_ai/training.py
index f629f61..f781f02 100644
--- a/bicleaner_ai/training.py
+++ b/bicleaner_ai/training.py
@@ -9,8 +9,10 @@ import random
 import fasttext
 
 try:
+    from . import __version__
     from .tokenizer import Tokenizer
 except (SystemError, ImportError):
+    from bicleaner_ai import __version__
     from tokenizer import Tokenizer
 
 # Porn removal classifier
@@ -270,6 +272,9 @@ def check_relative_path(path, filepath):
 def write_metadata(args, classifier, y_true, y_pred, lm_stats):
     out = args.metadata
 
+    # write current bicleaner ai version
+    out.write("bicleaner_ai_version: f{__version__}")
+
     precision = precision_score(y_true, y_pred)
     recall = recall_score(y_true, y_pred)
     f1 = f1_score(y_true, y_pred)