Fixes bert loading

racai-ai · Feb 24, 2024 · 7bd4deb · 7bd4deb
1 parent be5f83e
commit 7bd4deb
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/WebServiceModules/BERTAnnotator/annotator_api.py b/WebServiceModules/BERTAnnotator/annotator_api.py
@@ -3,6 +3,7 @@
 import argparse
 from flask import Flask, jsonify
 from annotator import NeuralAnnotator, BERTEntityTagger
+from bert import ReaderbenchSmall
 sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))
 from lib.saroj.gunicorn import StandaloneApplication
 from lib.saroj.input_data import get_input_data
@@ -51,7 +52,7 @@ def annotate_conllu():
     global tagger
 
     if tagger is None:
-        tagger = BERTEntityTagger(seq_len=256)
+        tagger = BERTEntityTagger(bert_model=ReaderbenchSmall())
         tagger.load(model_folder=args.MODEL)
     # end if