60 lines
1.4 KiB
Python
60 lines
1.4 KiB
Python
import argparse, json, sys, time, random
|
|
import spacy
|
|
from aitextgen import aitextgen
|
|
import string
|
|
|
|
def clean(text: str) -> str:
|
|
|
|
s = text.split('\n')
|
|
|
|
if(len(s) > 0):
|
|
tok_1 = s[0].split(' ')
|
|
if len(tok_1) > 0 and tok_1[0].strip() in string.punctuation:
|
|
s_1 = ' '.join(tok_1[1:])
|
|
s[0] = s_1.capitalize()
|
|
else:
|
|
s[0] = s[0].capitalize()
|
|
|
|
return '\n'.join(s)
|
|
|
|
def format(text: str) -> str:
|
|
|
|
return text.replace('\r\n', '\n').replace('\n\n', '\n')
|
|
|
|
def main() -> int:
|
|
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("-c", "--config", type=str, default="config.json", help="configuratin file")
|
|
p.add_argument("-i", "--iterations", type=int, default=10, help="number of iterations")
|
|
args = p.parse_args()
|
|
|
|
print(args)
|
|
|
|
with open(args.config) as f:
|
|
conf = json.load(f)
|
|
|
|
voices = []
|
|
for v in conf['voices']:
|
|
a = aitextgen(model_folder=v['model_dir'], tokenizer_file=v['tokeniser_file'])
|
|
voices.append({"name": v["name"].upper(), "a": a, "temp": float(v["temperature"])})
|
|
|
|
nbr_voices = len(voices)
|
|
current_voice = ""
|
|
for i in range(args.iterations):
|
|
rindex = random.randint(0, nbr_voices - 1)
|
|
v = voices[rindex]
|
|
if v['name'] != current_voice:
|
|
print("==========")
|
|
print(v['name'] + ":")
|
|
current_voice = v['name']
|
|
t = v['a'].generate(n=1, max_lenght=32, temperature=v['temp'], return_as_list=True)[0]
|
|
if t != None:
|
|
t = clean(t)
|
|
t = format(t)
|
|
print(t)
|
|
|
|
time.sleep(4)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main()) |