test 5 version

model.py

@@ -21,14 +21,16 @@ def preprocess(text):
 
 
 def predict(word_before, word_after):
-    prob_list = model[(word_before, word_after)].items()
+    prob_list = dict(Counter(model[(word_before, word_after)]).most_common(5)).items()
     predictions = []
-    total = 0.0
+    prob_sum = 0.0
     for key, value in prob_list:
-        total += value
+        prob_sum += value
         predictions.append(f'{key}:{value}')
-    if total == 0.0:
-        return 'the:1.0'
+    if prob_sum == 0.0:
+        return 'the:0:2 be:0.2 to:0.2 of:0.15 and:0.15 :0.1'
+    elif prob_sum < 1.0:
+        predictions.append(f':{1.0 - prob_sum}')
     return ' '.join(predictions)