1818
1919import pandas as pd
2020
21+
2122class WhatsAppChat :
2223 def __init__ (self , filename ):
2324 # Initialize the object by reading in the chat file, and counting word frequency and next words
@@ -44,8 +45,9 @@ def count_word_frequency(self):
4445 else :
4546 word_freq_dict [word ] = 1
4647 # Sort the word frequency dictionary by value in descending order and return it
47- sorted_word_freq_dict = dict (sorted (word_freq_dict .items (),
48- key = lambda x : - 1 * int (x [1 ])))
48+ sorted_word_freq_dict = dict (
49+ sorted (word_freq_dict .items (), key = lambda x : - 1 * int (x [1 ]))
50+ )
4951 return sorted_word_freq_dict
5052
5153 def find_next_words (self ):
@@ -57,7 +59,7 @@ def find_next_words(self):
5759 for i in range (len (words )):
5860 if i != len (words ) - 1 :
5961 current_word = words [i ]
60- next_word = words [i + 1 ].replace ("\n " , "" )
62+ next_word = words [i + 1 ].replace ("\n " , "" )
6163 if current_word in next_words_dict :
6264 if next_word in next_words_dict [current_word ]:
6365 next_words_dict [current_word ][next_word ] += 1
@@ -80,20 +82,21 @@ def predict_next_word(self, current_word):
8082 if current_word in self .chat .next_words_dict :
8183 # Get the dictionary of next words for the current word and sort it by frequency
8284 next_word_freq_dict = self .chat .next_words_dict [current_word ]
83- sorted_next_words = sorted (next_word_freq_dict . items (),
84- key = lambda x : x [1 ],
85- reverse = True )
85+ sorted_next_words = sorted (
86+ next_word_freq_dict . items (), key = lambda x : x [1 ], reverse = True
87+ )
8688 # Get the top num_predictions words and return them
87- next_words = [word [0 ] for word in sorted_next_words [:self .num_predictions ]]
89+ next_words = [word [0 ] for word in sorted_next_words [: self .num_predictions ]]
8890 return next_words
8991 else :
9092 # If the current word is not in the dictionary of next words, return None
9193 return None
9294
95+
9396# Create a WhatsAppChat object and a NextWordPredictor object after loading Chats.txt file
9497# set n as the length of words you want to predict after a given word
95- chat = WhatsAppChat (' Chats.txt' )
96- n , given_word = 3 , ' good'
98+ chat = WhatsAppChat (" Chats.txt" )
99+ n , given_word = 3 , " good"
97100predictor = NextWordPredictor (chat , num_predictions = n )
98101
99102# This prints the next n word predictions after the given_word
0 commit comments