NLP/bert-utils
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

修改predict代码

Browse Source
This commit is contained in:
joe 2019-01-30 09:57:02 +08:00
parent 3c4ed2ffe7
commit 0ca26576a4
1 changed files with 11 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
similarity.py
View File

@ -128,7 +128,7 @@ class BertSim:
def set_mode(self, mode):
self.mode = mode
self.estimator = self.get_estimator()
if mode == tf.estimator.ModeKeys.TRAIN:
if mode == tf.estimator.ModeKeys.PREDICT:
self.input_queue = Queue(maxsize=1)
self.output_queue = Queue(maxsize=1)
self.predict_thread = Thread(target=self.predict_from_queue, daemon=True)
@ -664,13 +664,13 @@ class BertSim:
if __name__ == '__main__':
sim = BertSim()
sim.set_mode(tf.estimator.ModeKeys.TRAIN)
sim.train()
sim.set_mode(tf.estimator.ModeKeys.EVAL)
sim.eval()
# sim.set_mode(tf.estimator.ModeKeys.PREDICT)
# while True:
# sentence1 = input('sentence1: ')
# sentence2 = input('sentence2: ')
# predict = sim.predict(sentence1, sentence2)
# print(f'similarity：{predict[0][1]}')
# sim.set_mode(tf.estimator.ModeKeys.TRAIN)
# sim.train()
# sim.set_mode(tf.estimator.ModeKeys.EVAL)
# sim.eval()
sim.set_mode(tf.estimator.ModeKeys.PREDICT)
while True:
sentence1 = input('sentence1: ')
sentence2 = input('sentence2: ')
predict = sim.predict(sentence1, sentence2)
print(f'similarity：{predict[0][1]}')