| 19 | |
| 20 | |
| 21 | class WhatsAppChat: |
| 22 | def __init__(self, filename): |
| 23 | # Initialize the object by reading in the chat file, and counting word frequency and next words |
| 24 | self.titles = self.read_chat_file(filename) |
| 25 | self.word_freq_dict = self.count_word_frequency() |
| 26 | self.next_words_dict = self.find_next_words() |
| 27 | |
| 28 | def read_chat_file(self, filename): |
| 29 | """Reads in a WhatsApp chat text file and returns a list of titles.""" |
| 30 | chat_df = pd.read_fwf(filename, header=None) |
| 31 | # Extract the chat titles from the DataFrame and return them as a list |
| 32 | titles = [title for title in chat_df[2]] |
| 33 | return titles |
| 34 | |
| 35 | def count_word_frequency(self): |
| 36 | """Counts the frequency of each word in a list of titles.""" |
| 37 | word_freq_dict = {} |
| 38 | for title in self.titles: |
| 39 | # Split each title into words and count their frequency |
| 40 | for word in str(title).split(): |
| 41 | if len(word) > 1: |
| 42 | if word in word_freq_dict.keys(): |
| 43 | word_freq_dict[word] += 1 |
| 44 | else: |
| 45 | word_freq_dict[word] = 1 |
| 46 | # Sort the word frequency dictionary by value in descending order and return it |
| 47 | sorted_word_freq_dict = dict( |
| 48 | sorted(word_freq_dict.items(), key=lambda x: -1 * int(x[1])) |
| 49 | ) |
| 50 | return sorted_word_freq_dict |
| 51 | |
| 52 | def find_next_words(self): |
| 53 | """Finds the most common next word for each word in a list of titles.""" |
| 54 | next_words_dict = {} |
| 55 | for title in self.titles: |
| 56 | # Split each title into words and count the frequency of each next word for each word |
| 57 | words = str(title).split() |
| 58 | for i in range(len(words)): |
| 59 | if i != len(words) - 1: |
| 60 | current_word = words[i] |
| 61 | next_word = words[i + 1].replace("\n", "") |
| 62 | if current_word in next_words_dict: |
| 63 | if next_word in next_words_dict[current_word]: |
| 64 | next_words_dict[current_word][next_word] += 1 |
| 65 | else: |
| 66 | next_words_dict[current_word][next_word] = 1 |
| 67 | else: |
| 68 | next_words_dict[current_word] = {next_word: 1} |
| 69 | # Return the dictionary of next words |
| 70 | return next_words_dict |
| 71 | |
| 72 | |
| 73 | class NextWordPredictor: |