File size: 1,263 Bytes
d0ffe9c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import glob
import os
import random
import re
wild_card_regex = r'(\A|\W)__([\w-]+)__(\W|\Z)'
def create_wild_card_map(wild_card_dir):
result = {}
if os.path.isdir(wild_card_dir):
txt_list = glob.glob( os.path.join(wild_card_dir ,"**/*.txt"), recursive=True)
for txt in txt_list:
basename_without_ext = os.path.splitext(os.path.basename(txt))[0]
with open(txt, encoding='utf-8') as f:
try:
result[basename_without_ext] = [s.rstrip() for s in f.readlines()]
except Exception as e:
print(e)
print("can not read ", txt)
return result
def replace_wild_card_token(match_obj, wild_card_map):
m1 = match_obj.group(1)
m3 = match_obj.group(3)
dict_name = match_obj.group(2)
if dict_name in wild_card_map:
token_list = wild_card_map[dict_name]
token = token_list[random.randint(0,len(token_list)-1)]
return m1+token+m3
else:
return match_obj.group(0)
def replace_wild_card(prompt, wild_card_dir):
wild_card_map = create_wild_card_map(wild_card_dir)
prompt = re.sub(wild_card_regex, lambda x: replace_wild_card_token(x, wild_card_map ), prompt)
return prompt
|