From 307c291862010cc662b4d05c7352aac2ca4b0422 Mon Sep 17 00:00:00 2001
From: Sebastian Kutny <sebkutny@gmail.com>
Date: Tue, 28 Mar 2023 17:25:17 +0200
Subject: [PATCH] Added rhymes in lyrics generation.

---
 main.py         |  8 +++++---
 markov_model.py | 30 +++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/main.py b/main.py
index e93570b..18abafc 100644
--- a/main.py
+++ b/main.py
@@ -26,13 +26,15 @@ def generate_song(name):
     n_gram = int(input("Select number of words in Markov state: "))
     number_of_verses = int(input("Select number of verses: "))
     words_in_verses = int((int(input("Select number of words in verses: ")) - 1) / n_gram)
-    degree_of_chain = int(input("Select degree of chain: "))
-    model = create_markov_model(dataset, n_gram, degree_of_chain)
+    # degree_of_chain = int(input("Select degree of chain: "))
+    model = create_markov_model(dataset, n_gram)
     print('\n')
     last_state = random.choice(list(model.keys()))
+    rime = None
     for i in range(number_of_verses):
-        generated_lyrics, last_state = generate_lyrics(model, last_state, words_in_verses)
+        generated_lyrics, last_state = generate_lyrics(model, last_state, words_in_verses, True if i == 0 else False, rime)
         print(generated_lyrics)
+        rime = last_state
         last_state = random.choices(list(model[last_state].keys()),
                                     list(model[last_state].values()))[0]
 
diff --git a/markov_model.py b/markov_model.py
index f426a6c..55df1f1 100644
--- a/markov_model.py
+++ b/markov_model.py
@@ -1,5 +1,6 @@
 import random
 import re
+from nltk import SyllableTokenizer
 from nltk.tokenize import word_tokenize
 import pandas as pd
 import numpy as np
@@ -29,7 +30,7 @@ def clean_data(name):
     return dataset
 
 
-def create_markov_model(dataset, n_gram, n_step):
+def create_markov_model(dataset, n_gram):
     markov_model = {}
     for i in range(len(dataset) - 1 - 2 * n_gram):
         current_state, next_state = "", ""
@@ -73,16 +74,35 @@ def create_markov_model(dataset, n_gram, n_step):
     return markov_model
 
 
-def generate_lyrics(markov_model, start, limit):
+def generate_lyrics(markov_model, start, limit, isStartingVerse, rime):
     n = 0
     current_state = start
     lyrics = ""
     lyrics += current_state + " "
     lyrics = lyrics[0].upper() + lyrics[1:]
     while n < limit:
-        next_state = random.choices(list(markov_model[current_state].keys()),
-                                    list(markov_model[current_state].values()))
-        current_state = next_state[0]
+        if n == limit - 1 and not isStartingVerse:
+            rime = rime.split(" ")[-1]
+            tk = SyllableTokenizer()
+            rime_syllab = tk.tokenize(rime)[-1]
+            rime_states = {}
+            for state, probability in markov_model[current_state].items():
+                word = state.split(" ")[-1]
+                syllab = tk.tokenize(word)[-1]
+                if rime_syllab == syllab and rime != word:
+                    rime_states.update({state: probability})
+            if rime_states:
+                next_state = random.choices(list(rime_states.keys()),
+                                            list(rime_states.values()))
+                current_state = next_state[0]
+            else:
+                next_state = random.choices(list(markov_model[current_state].keys()),
+                                            list(markov_model[current_state].values()))
+                current_state = next_state[0]
+        else:
+            next_state = random.choices(list(markov_model[current_state].keys()),
+                                        list(markov_model[current_state].values()))
+            current_state = next_state[0]
         lyrics += current_state + " "
         n += 1
     return lyrics, current_state