using beam size

bedapudi6788 · Feb 1, 2019 · ad1fcc6 · ad1fcc6
1 parent e88d8a5
commit ad1fcc6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/deepcorrect/deepcorrect.py b/deepcorrect/deepcorrect.py
@@ -8,11 +8,11 @@ def __init__(self, params_path, checkpoint_path):
         DeepCorrect.deepcorrect_model = build_model(params_path)
         DeepCorrect.deepcorrect_model[0].load_weights(checkpoint_path)
 
-    def correct(self, sentence):
+    def correct(self, sentence, beam_size = 1):
         if not DeepCorrect.deepcorrect_model:
             print('Please load the model first')
 
         sentence = sentence.strip()
-        sentence = infer(sentence, DeepCorrect.deepcorrect_model[0], DeepCorrect.deepcorrect_model[1])
+        sentence = infer(sentence, DeepCorrect.deepcorrect_model[0], DeepCorrect.deepcorrect_model[1], beam_size = beam_size)
 
         return sentence