| 71 | |
| 72 | |
| 73 | class NextWordPredictor: |
| 74 | def __init__(self, chat, num_predictions=1): |
| 75 | # Initialize the object with a WhatsAppChat object and the number of predictions to make |
| 76 | self.chat = chat |
| 77 | self.num_predictions = num_predictions |
| 78 | |
| 79 | def predict_next_word(self, current_word): |
| 80 | """Predicts the next word given a current word and a dictionary of next words.""" |
| 81 | if current_word in self.chat.next_words_dict: |
| 82 | # Get the dictionary of next words for the current word and sort it by frequency |
| 83 | next_word_freq_dict = self.chat.next_words_dict[current_word] |
| 84 | sorted_next_words = sorted( |
| 85 | next_word_freq_dict.items(), key=lambda x: x[1], reverse=True |
| 86 | ) |
| 87 | # Get the top num_predictions words and return them |
| 88 | next_words = [word[0] for word in sorted_next_words[: self.num_predictions]] |
| 89 | return next_words |
| 90 | else: |
| 91 | # If the current word is not in the dictionary of next words, return None |
| 92 | return None |
| 93 | |
| 94 | |
| 95 | # Create a WhatsAppChat object and a NextWordPredictor object after loading Chats.txt file |