mirror of
https://github.com/brandons209/Red-bot-Cogs.git
synced 2024-06-26 18:10:26 +12:00
overhauled structure, saving data, and integrated economy cog
This commit is contained in:
brandons209
parent
f8ac9cfe69
commit
b05034cb07
|
|
@ -12,11 +12,13 @@ import pickle
|
|||
import time
|
||||
import asyncio
|
||||
|
||||
from .utils.dataIO import dataIO
|
||||
from .utils.chat_formatting import warning
|
||||
|
||||
#loads dictionary from file
|
||||
def _load_dict(path):
|
||||
with open(path, 'rb') as file:
|
||||
dict = pickle.load(file)
|
||||
dict = pickle.load(file)
|
||||
return dict
|
||||
|
||||
#dictionaries for tokenizing puncuation and converting it back
|
||||
|
|
@ -59,15 +61,35 @@ def _sample(prediction, temp=0):
|
|||
probabilities = np.random.multinomial(1, prediction, 1)
|
||||
return np.argmax(probabilities)
|
||||
|
||||
def _resolve_role_list(server: discord.Server, roles: list) -> list:
|
||||
gen = (_role_from_string(server, name) for name in roles)
|
||||
return list(filter(None, gen))
|
||||
|
||||
def _role_from_string(server, rolename, roles=None):
|
||||
if roles is None:
|
||||
roles = server.roles
|
||||
|
||||
roles = [r for r in roles if r is not None]
|
||||
role = discord.utils.find(lambda r: r.name.lower() == rolename.lower(), roles)
|
||||
# if couldnt find by role name, try to find by role id
|
||||
if role is None:
|
||||
role = discord.utils.find(lambda r: r.id == rolename, roles)
|
||||
|
||||
return role
|
||||
|
||||
def format_list(*items, join='and', delim=', '):
|
||||
if len(items) > 1:
|
||||
return (' %s ' % join).join((delim.join(items[:-1]), items[-1]))
|
||||
elif items:
|
||||
return items[0]
|
||||
else:
|
||||
return ''
|
||||
|
||||
"""This cog generates scripts based on imported model, I used a keras model. """
|
||||
class ScriptCog:
|
||||
|
||||
def __init__(self, bot):
|
||||
self.bot = bot
|
||||
|
||||
os.makedirs("data/scriptcog/", exist_ok=True)
|
||||
os.makedirs("data/scriptcog/dicts", exist_ok=True)
|
||||
self.model_path = "data/scriptcog/model.h5"
|
||||
self.dict_path = "data/scriptcog/dicts/"
|
||||
|
||||
|
|
@ -78,19 +100,13 @@ class ScriptCog:
|
|||
except:
|
||||
self.model = None
|
||||
|
||||
if os.path.isfile("data/scriptcog/config.txt"):
|
||||
with open("data/scriptcog/config.txt", "r") as f:
|
||||
content = f.read()
|
||||
content = content.split()
|
||||
self.word_limit = int(content[0])
|
||||
self.cooldown_limit = int(content[1])
|
||||
self.tv_show = content[2]
|
||||
else:
|
||||
with open("data/scriptcog/config.txt", "w") as f:
|
||||
f.write("100 30 MLP")
|
||||
self.word_limit = 100
|
||||
self.cooldown_limit = 30
|
||||
self.tv_show = "MLP"
|
||||
self.settings_path = "data/scriptcog/settings.json"
|
||||
self.settings = dataIO.load_json(self.settings_path)
|
||||
|
||||
self.default_word_limit = 300
|
||||
self.default_cooldown_limit = 30
|
||||
self.default_tv_show = "My Little Pony"
|
||||
self.default_price = 0
|
||||
|
||||
try:
|
||||
self.word_to_int = _load_dict(self.dict_path + 'word_to_int.pkl')
|
||||
|
|
@ -101,41 +117,178 @@ class ScriptCog:
|
|||
self.int_to_word = None
|
||||
self.sequence_length = None
|
||||
|
||||
def _write_config(self):
|
||||
with open("data/scriptcog/config.txt", "w") as f:
|
||||
f.write("{} {} {}".format(self.word_limit, self.cooldown_limit, self.tv_show))
|
||||
@commands.group(pass_context=True, invoke_without_command=True, no_pm=True)
|
||||
@checks.admin_or_permissions(administrator=True)
|
||||
async def genscriptset(self, ctx):
|
||||
if ctx.invoked_subcommand is None:
|
||||
await self.bot.send_cmd_help(ctx)
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
@genscriptset.command(pass_context=True, no_pm=True, name="word-limit")
|
||||
@checks.is_owner()
|
||||
async def setwordlimit(self, ctx, num_words : int = 100):
|
||||
#if ctx.invoked_subcommand is None:
|
||||
# await self.bot.say("Usage: setwordlimit limit")
|
||||
# return
|
||||
self.word_limit = num_words
|
||||
self._write_config()
|
||||
await self.bot.say("Maximum number of words is now {}".format(self.word_limit))
|
||||
async def genscriptset_set_word_limit(self, ctx, num_words : int):
|
||||
"""
|
||||
Set the word limit for generating scripts.
|
||||
"""
|
||||
try:
|
||||
self.settings[ctx.message.server.id]["WORD_LIMIT"] = num_words
|
||||
except:
|
||||
self.settings[ctx.message.server.id] = {}
|
||||
self.settings[ctx.message.server.id]["WORD_LIMIT"] = num_words
|
||||
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
await self.bot.say("Maximum number of words is now {}".format(num_words))
|
||||
|
||||
@genscriptset.command(pass_context=True, no_pm=True, name="tv-show")
|
||||
@checks.is_owner()
|
||||
async def genscriptset_set_tv_show(self, ctx, *, show):
|
||||
"""
|
||||
Sets the TV show the scripts are generated from.
|
||||
"""
|
||||
try:
|
||||
self.settings[ctx.message.server.id]["TV_SHOW"] = show
|
||||
except:
|
||||
self.settings[ctx.message.server.id] = {}
|
||||
self.settings[ctx.message.server.id]["TV_SHOW"] = show
|
||||
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
await self.bot.say("TV show is now {}.".format(show))
|
||||
|
||||
@genscriptset.command(pass_context=True, no_pm=True, name="cooldown")
|
||||
@checks.is_owner()
|
||||
async def genscriptset_set_cooldown(self, ctx, cooldown : int):
|
||||
"""
|
||||
Sets the cooldown period between generating scripts in seconds.
|
||||
"""
|
||||
try:
|
||||
self.settings[ctx.message.server.id]["COOLDOWN"] = cooldown
|
||||
except:
|
||||
self.settings[ctx.message.server.id] = {}
|
||||
self.settings[ctx.message.server.id]["COOLDOWN"] = cooldown
|
||||
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
await self.bot.say("Script cooldown is now {}.".format(cooldown))
|
||||
|
||||
@genscriptset.command(pass_context=True, no_pm=True, name="price")
|
||||
async def genscriptset_set_price(self, ctx, price : int):
|
||||
"""
|
||||
Sets the price for generating scripts.
|
||||
"""
|
||||
try:
|
||||
self.settings[ctx.message.server.id]["PRICE"] = price
|
||||
except:
|
||||
self.settings[ctx.message.server.id] = {}
|
||||
self.settings[ctx.message.server.id]["PRICE"] = price
|
||||
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
await self.bot.say("Price is now {}.".format(price))
|
||||
|
||||
@genscriptset.command(pass_context=True, no_pm=True, name="free-roles")
|
||||
async def genscriptset_free_roles(self, ctx, *, rolelist=None):
|
||||
"""Set roles that do not have to pay to generate scripts.
|
||||
|
||||
COMMA SEPARATED LIST (e.g. Admin,Staff,Mod), Can also use role IDs as well.
|
||||
|
||||
To get current list, run command with no roles.
|
||||
|
||||
Add role_list_clear as the role to clear the server\'s free role list.
|
||||
"""
|
||||
server = ctx.message.server
|
||||
current_roles = _resolve_role_list(server, self.settings[server.id].get("FREE_ROLE_LIST", []))
|
||||
|
||||
if rolelist is None:
|
||||
if current_roles:
|
||||
names_list = format_list(*(r.name for r in current_roles))
|
||||
await self.bot.say("Current list of roles that do not have to pay: {}".format(names_list))
|
||||
else:
|
||||
await self.bot.say("No roles defined.")
|
||||
return
|
||||
elif "role_list_clear" in rolelist.lower():
|
||||
await self.bot.say("Free role list cleared.")
|
||||
self.settings[server.id]["FREE_ROLE_LIST"] = []
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
return
|
||||
|
||||
found_roles = set()
|
||||
notfound_names = set()
|
||||
|
||||
for lookup in rolelist.split(","):
|
||||
lookup = lookup.strip()
|
||||
role = _role_from_string(server, lookup)
|
||||
|
||||
if role:
|
||||
found_roles.add(role)
|
||||
else:
|
||||
notfound_names.add(lookup)
|
||||
|
||||
if notfound_names:
|
||||
fmt_list = format_list(*("`{}`".format(x) for x in notfound_names))
|
||||
await self.bot.say(warning("These roles were not found: {}\n\nPlease try again.".format(fmt_list)))
|
||||
elif server.default_role in found_roles:
|
||||
await self.bot.say(warning("The everyone role cannot be added.\n\nPlease try again."))
|
||||
elif found_roles == set(current_roles):
|
||||
await self.bot.say("No changes to make.")
|
||||
else:
|
||||
if server.id not in self.settings:
|
||||
self.settings[server.id] = {}
|
||||
else:
|
||||
extra = ""
|
||||
|
||||
self.settings[server.id]["FREE_ROLE_LIST"] = [r.id for r in found_roles]
|
||||
dataIO.save_json(self.settings_path, self.settings)
|
||||
|
||||
fmt_list = format_list(*(r.name for r in found_roles))
|
||||
await self.bot.say("These roles will not have to pay for scripts: {}.{}".format(fmt_list, extra))
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
async def genscriptinfo(self, ctx):
|
||||
await self.bot.say("Word Limit: {}, Cooldown Time: {}, Show: {}".format(self.word_limit, self.cooldown_limit, self.tv_show))
|
||||
server_id = ctx.message.server.id
|
||||
word_limit = self.settings.get(server_id, {}).get("WORD_LIMIT", self.default_word_limit)
|
||||
cooldown = self.settings.get(server_id, {}).get("COOLDOWN", self.default_cooldown_limit)
|
||||
tv_show = self.settings.get(server_id, {}).get("TV_SHOW", self.default_tv_show)
|
||||
price = self.settings.get(server_id, {}).get("PRICE", self.default_price)
|
||||
await self.bot.say("Word Limit: {}, Cooldown Time: {}, Show: {}, Price: {}".format(word_limit, cooldown, tv_show, price))
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
async def genscripthelp(self, ctx):
|
||||
await self.bot.say("--------------------\nGenerate original TV scripts for {} using Neural Networks!\nUsage: `genscript number_of_words_to_generate word_variance starting_text`\nUse starting texts such as:\n`\"pinkie pie::\"`\n`fluttershy::`\n`\"twilight sparkle::\"`\nor other names of characters in the show. Otherwise, you can use any words said in the show.\n\n**If your starting text has spaces, enclose it with quotes: \"\"**\nWord variance helps gives the script better results. A variance of 0 will mean that with the same starting text, it will always have the same output. Variance up to 1.0 will give more variety to words, however going closer to 1 can introduce more grammar and spelling mistakes.\n-------------------".format(self.tv_show))
|
||||
await self.bot.say("--------------------\nGenerate original TV scripts for {} using Neural Networks!\nUsage: `genscript <number of words to generate> <word variance> <starting text>`\nUse starting texts such as:\n`pinkie pie::`\n`fluttershy::`\n`twilight sparkle::`\nor other names of characters in the show. Otherwise, you can use any words said in the show.\n\nWord variance helps gives the script better results. A variance of 0 will mean that with the same starting text, it will always have the same output. Variance up to 1.0 will give more variety to words, however going closer to 1 can introduce more grammar and spelling mistakes.\n-------------------".format(self.settings.get(ctx.message.server.id, {}).get("TV_SHOW", self.default_tv_show)))
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
@checks.is_owner()
|
||||
async def settvshow(self, ctx, show : str = "MLP"):
|
||||
self.tv_show = show
|
||||
self._write_config()
|
||||
await self.bot.say("TV show is now {}.".format(self.tv_show))
|
||||
async def genscript(self, ctx, num_words_to_generate : int, variance : float, *, seed):
|
||||
"""
|
||||
Generate a script using the power of Neural Networks!
|
||||
Please use the genscripthelp command to get a complete explaination of how to use the command.
|
||||
"""
|
||||
server_id = ctx.message.server.id
|
||||
word_limit = self.settings.get(server_id, {}).get("WORD_LIMIT", self.default_word_limit)
|
||||
cooldown = self.settings.get(server_id, {}).get("COOLDOWN", self.default_cooldown_limit)
|
||||
price = self.settings.get(server_id, {}).get("PRICE", self.default_price)
|
||||
user = ctx.message.author
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
@checks.is_owner()
|
||||
async def setcooldown(self, ctx, cooldown : int = 30):
|
||||
self.cooldown_limit = cooldown
|
||||
self._write_config()
|
||||
await self.bot.say("Script cooldown is now {}.".format(self.cooldown_limit))
|
||||
if num_words_to_generate > word_limit:
|
||||
await self.bot.say("Please keep script sizes to {} words or less.".format(word_limit))
|
||||
return
|
||||
elif time.time() - self.cooldown < cooldown:
|
||||
await self.bot.say("Sorry, I am cooling down, please wait {:.0f} seconds.".format(cooldown - (time.time() - self.cooldown)))
|
||||
return
|
||||
|
||||
if not self.check_free_role(user):
|
||||
return_val = self.charge_user(user, price)
|
||||
if return_val == 1: #if this worked, dont need to check when getting econ cog
|
||||
balance = self.bot.get_cog('Economy').bank.get_balance(user)
|
||||
await self.bot.say("Charged: {}, Balance: {}".format(price, balance))
|
||||
else:
|
||||
await self.charge_user_check(return_val, price, user)
|
||||
return
|
||||
|
||||
self.cooldown = time.time()
|
||||
|
||||
if variance > 1.0:
|
||||
variance = 1.0
|
||||
elif variance < 0:
|
||||
variance = 0
|
||||
|
||||
await self.bot.say("Generating script, please wait...")
|
||||
await self.get_model_output(num_words_to_generate, variance, seed)
|
||||
|
||||
async def get_model_output(self, num_words, temp, seed):
|
||||
input_text = seed
|
||||
|
|
@ -162,27 +315,58 @@ class ScriptCog:
|
|||
await self.bot.say(result)
|
||||
await self.bot.say("------------------------")
|
||||
|
||||
@commands.command(pass_context=True, no_pm=True)
|
||||
async def genscript(self, ctx, num_words_to_generate : int = 100, variance : float = 0.5, seed : str = "pinkie pie::"):
|
||||
#if ctx.invoked_subcommand is None:
|
||||
# await self.bot.say("Usage: genscript num_words randomness(between 0 and 1) seed_text")
|
||||
# return
|
||||
if num_words_to_generate > self.word_limit:
|
||||
await self.bot.say("Please keep script sizes to {} words or less.".format(self.word_limit))
|
||||
return
|
||||
elif time.time() - self.cooldown < self.cooldown_limit:
|
||||
await self.bot.say("Sorry, I am cooling down, please wait {:.0f} seconds.".format(self.cooldown_limit - (time.time() - self.cooldown)))
|
||||
def charge_user(self, user, amount):
|
||||
"""
|
||||
Takes a user and a amount and charges the user. Returns 1 on success, 0 if economy cog cannot be loaded, -1 if the user does not have a bank account, and -2 if the user doesn't have enough credits.
|
||||
"""
|
||||
econ_cog = self.bot.get_cog('Economy')
|
||||
if not econ_cog:
|
||||
return 0
|
||||
if not econ_cog.bank.account_exists(user):
|
||||
return -1
|
||||
if not econ_cog.bank.can_spend(user, amount):
|
||||
return -2
|
||||
econ_cog.bank.withdraw_credits(user, amount)
|
||||
return 1
|
||||
|
||||
async def charge_user_check(self, return_val, amount, user):
|
||||
"""
|
||||
takes in return vals from charge_user and gives apporiate response
|
||||
"""
|
||||
econ_cog = self.bot.get_cog('Economy')
|
||||
if not econ_cog:
|
||||
await self.bot.say("Error loading economy cog!")
|
||||
return
|
||||
if return_val == 0:
|
||||
await self.bot.say("Economy cog not found! Please check to make sure the economy cog is loaded.")
|
||||
elif return_val == -1:
|
||||
await self.bot.say("You appear to not have a bank account. Use [p]bank register to open an account.")
|
||||
elif return_val == -2:
|
||||
await self.bot.say("You do not have enough credits. The sound command costs {} and you have {} credits in your bank account.".format(amount, econ_cog.bank.get_balance(user)))
|
||||
|
||||
self.cooldown = time.time()
|
||||
def check_free_role(self, user):
|
||||
"""
|
||||
Checks if the user has a role that grants them free sounds. Returns 1 if the user does have one of these roles, zero otherwise.
|
||||
"""
|
||||
server = user.server
|
||||
free_roles = self.settings.get(server.id, {}).get("FREE_ROLE_LIST", [])
|
||||
if not free_roles:
|
||||
return 0
|
||||
user_roles = [r.id for r in user.roles]
|
||||
|
||||
if variance > 1.0:
|
||||
variance = 1.0
|
||||
elif variance < 0:
|
||||
variance = 0
|
||||
for role in free_roles:
|
||||
if role in user_roles:
|
||||
return 1
|
||||
|
||||
await self.bot.say("Generating script, please wait...")
|
||||
await self.get_model_output(num_words_to_generate, variance, seed)
|
||||
return 0
|
||||
|
||||
def check_folders():
|
||||
os.makedirs("data/scriptcog", exist_ok=True)
|
||||
os.makedirs("data/scriptcog/dicts", exist_ok=True)
|
||||
f = "data/scriptcog/settings.json"
|
||||
if not dataIO.is_valid_json(f):
|
||||
dataIO.save_json(f, {})
|
||||
|
||||
def setup(bot):
|
||||
check_folders()
|
||||
bot.add_cog(ScriptCog(bot))
|
||||
|
|
|
|||
Loading…
Reference in a new issue