choose_words.py

import numpy as np
from nltk.tokenize import word_tokenize
from nltk import pos_tag

def pick_syl(target):
    si = 0
    choices = list()
    while si != target:
        if target-si != 1:
            cr = np.random.randint(1, target-si)
            if cr > 3: cr = 1
            choices.append(cr)
            si += cr
        else: 
            choices.append(1)
            return choices
    return choices

def syl_of_size(sdict, size):
    return [w for w in sdict.keys() if sdict[w] == size]

def choose_words(sdict, nr_syl, line):
    words = list()
    ix = 0
    for sw in pick_syl(nr_syl):
        if ix == 0 and line == 0:
            sdict1 = dict([(sk, sv) for sk, sv in sdict.items() if pos_tag(word_tokenize(sk), tagset='universal')[0][1]=='NOUN'])
            gsyls = syl_of_size(sdict1, sw)
        elif ix == 1 and line == 0:
            sdict2 = dict([(sk, sv) for sk, sv in sdict.items() if pos_tag(word_tokenize(sk), tagset='universal')[0][1]=='VERB'])
            gsyls = syl_of_size(sdict2, sw)
        elif ix == 2 and line == 0:
            sdict3 = dict([(sk, sv) for sk, sv in sdict.items() if pos_tag(word_tokenize(sk), tagset='universal')[0][1] not in ['NOUN', 'VERB']])
            gsyls = syl_of_size(sdict3, sw)
        else:
            gsyls = syl_of_size(sdict, sw)
        cwr = np.random.randint(len(gsyls))
        chosen = gsyls[cwr]
        words.append(chosen)
        ix += 1
    return words