diff --git a/README.md b/README.md index 980e762..c37f84f 100644 --- a/README.md +++ b/README.md @@ -21,7 +21,7 @@ Cog Function | infochannel | **Beta** |
Create a channel to display server infoJust released, please report bugs
| | lovecalculator | **Alpha** |
Calculate the love between two users[Snap-Ons] Just updated to V3
| | lseen | **Alpha** |
Track when a member was last onlineAlpha release, please report bugs
| -| nudity | **Incomplete** |
Checks for NSFW images posted in non-NSFW channelsLibrary this is based on has a bug, waiting for author to merge my PR
| +| nudity | **Alpha** |
Checks for NSFW images posted in non-NSFW channelsSwitched libraries, now functional
| | planttycoon | **Alpha** |
Grow your own plants![Snap-Ons] Updated to V3, likely to contain bugs
| | qrinvite | **Alpha** |
Create a QR code invite for the serverAlpha release, please report any bugs
| | reactrestrict | **Alpha** |
Removes reactions by role per channelA bit clunky, but functional
| diff --git a/chatter/README.md b/chatter/README.md index 933162a..e8c03d6 100644 --- a/chatter/README.md +++ b/chatter/README.md @@ -162,12 +162,53 @@ This command trains Chatter on the specified channel based on the configured settings. This can take a long time to process. +### Train Ubuntu + +``` +[p]chatter trainubuntu +``` +*WARNING:* This will trigger a large download and use a lot of processing power + +This command trains Chatter on the publicly available Ubuntu Dialogue Corpus. (It'll talk like a geek) + + ## Switching Algorithms ``` [p]chatter algorithm X ``` +or +``` +[p]chatter algo X 0.95 +``` Chatter can be configured to use one of three different Similarity algorithms. Changing this can help if the response speed is too slow, but can reduce the accuracy of results. + +The second argument is the maximum similarity threshold, +lowering that will make the bot stop searching sooner. + +Default maximum similarity threshold is 0.90 + + +## Switching Pretrained Models + +``` +[p]chatter model X +``` + +Chatter can be configured to use one of three pretrained statistical models for English. + +I have not noticed any advantage to changing this, +but supposedly it would help by splitting the search term into more useful parts. + +See [here](https://spacy.io/models) for more info on spaCy models. + +Before you're able to use the *large* model (option 3), you must install it through pip. + +*Warning:* This is ~800MB download. + +``` +[p]pipinstall https://github.com/explosion/spacy-models/releases/download/en_core_web_lg-2.3.1/en_core_web_lg-2.3.1.tar.gz#egg=en_core_web_lg +``` diff --git a/chatter/chat.py b/chatter/chat.py index f7e8944..76ee56c 100644 --- a/chatter/chat.py +++ b/chatter/chat.py @@ -1,4 +1,5 @@ import asyncio +import logging import os import pathlib from datetime import datetime, timedelta @@ -7,13 +8,16 @@ import discord from chatterbot import ChatBot from chatterbot.comparisons import JaccardSimilarity, LevenshteinDistance, SpacySimilarity from chatterbot.response_selection import get_random_response -from chatterbot.trainers import ChatterBotCorpusTrainer, ListTrainer +from chatterbot.trainers import ChatterBotCorpusTrainer, ListTrainer, UbuntuCorpusTrainer from redbot.core import Config, commands from redbot.core.commands import Cog from redbot.core.data_manager import cog_data_path +from redbot.core.utils.predicates import MessagePredicate +log = logging.getLogger("red.fox_v3.chat") -class ENG_LG: # TODO: Add option to use this large model + +class ENG_LG: ISO_639_1 = "en_core_web_lg" ISO_639 = "eng" ENGLISH_NAME = "English" @@ -25,6 +29,12 @@ class ENG_MD: ENGLISH_NAME = "English" +class ENG_SM: + ISO_639_1 = "en_core_web_sm" + ISO_639 = "eng" + ENGLISH_NAME = "English" + + class Chatter(Cog): """ This cog trains a chatbot that will talk like members of your Guild @@ -39,7 +49,13 @@ class Chatter(Cog): path: pathlib.Path = cog_data_path(self) self.data_path = path / "database.sqlite3" - self.chatbot = self._create_chatbot(self.data_path, SpacySimilarity, 0.45, ENG_MD) + # TODO: Move training_model and similarity_algo to config + # TODO: Add an option to see current settings + + self.tagger_language = ENG_MD + self.similarity_algo = SpacySimilarity + self.similarity_threshold = 0.90 + self.chatbot = self._create_chatbot() # self.chatbot.set_trainer(ListTrainer) # self.trainer = ListTrainer(self.chatbot) @@ -49,18 +65,18 @@ class Chatter(Cog): self.loop = asyncio.get_event_loop() - def _create_chatbot( - self, data_path, similarity_algorithm, similarity_threshold, tagger_language - ): + def _create_chatbot(self): + return ChatBot( "ChatterBot", storage_adapter="chatterbot.storage.SQLStorageAdapter", - database_uri="sqlite:///" + str(data_path), - statement_comparison_function=similarity_algorithm, + database_uri="sqlite:///" + str(self.data_path), + statement_comparison_function=self.similarity_algo, response_selection_method=get_random_response, logic_adapters=["chatterbot.logic.BestMatch"], - # maximum_similarity_threshold=similarity_threshold, - tagger_language=tagger_language, + maximum_similarity_threshold=self.similarity_threshold, + tagger_language=self.tagger_language, + logger=log, ) async def _get_conversation(self, ctx, in_channel: discord.TextChannel = None): @@ -130,6 +146,11 @@ class Chatter(Cog): return out + def _train_ubuntu(self): + trainer = UbuntuCorpusTrainer(self.chatbot) + trainer.train() + return True + def _train_english(self): trainer = ChatterBotCorpusTrainer(self.chatbot) # try: @@ -182,14 +203,18 @@ class Chatter(Cog): try: os.remove(self.data_path) except PermissionError: - await ctx.maybe_send_embed("Failed to clear training database. Please wait a bit and try again") + await ctx.maybe_send_embed( + "Failed to clear training database. Please wait a bit and try again" + ) - self._create_chatbot(self.data_path, SpacySimilarity, 0.45, ENG_MD) + self._create_chatbot() await ctx.tick() - @chatter.command(name="algorithm") - async def chatter_algorithm(self, ctx: commands.Context, algo_number: int): + @chatter.command(name="algorithm", aliases=["algo"]) + async def chatter_algorithm( + self, ctx: commands.Context, algo_number: int, threshold: float = None + ): """ Switch the active logic algorithm to one of the three. Default after reload is Spacy @@ -198,17 +223,61 @@ class Chatter(Cog): 2: Levenshtein """ - algos = [(SpacySimilarity, 0.45), (JaccardSimilarity, 0.75), (LevenshteinDistance, 0.75)] + algos = [SpacySimilarity, JaccardSimilarity, LevenshteinDistance] if algo_number < 0 or algo_number > 2: await ctx.send_help() return - self.chatbot = self._create_chatbot( - self.data_path, algos[algo_number][0], algos[algo_number][1], ENG_MD - ) + if threshold is not None: + if threshold >= 1 or threshold <= 0: + await ctx.maybe_send_embed( + "Threshold must be a number between 0 and 1 (exclusive)" + ) + return + else: + self.similarity_algo = threshold - await ctx.tick() + self.similarity_algo = algos[algo_number] + async with ctx.typing(): + self.chatbot = self._create_chatbot() + + await ctx.tick() + + @chatter.command(name="model") + async def chatter_model(self, ctx: commands.Context, model_number: int): + """ + Switch the active model to one of the three. Default after reload is Medium + + 0: Small + 1: Medium + 2: Large (Requires additional setup) + """ + + models = [ENG_SM, ENG_MD, ENG_LG] + + if model_number < 0 or model_number > 2: + await ctx.send_help() + return + + if model_number == 2: + await ctx.maybe_send_embed( + "Additional requirements needed. See guide before continuing.\n" "Continue?" + ) + pred = MessagePredicate.yes_or_no(ctx) + try: + await self.bot.wait_for("message", check=pred, timeout=30) + except TimeoutError: + await ctx.send("Response timed out, please try again later.") + return + if not pred.result: + return + + self.tagger_language = models[model_number] + async with ctx.typing(): + self.chatbot = self._create_chatbot() + + await ctx.maybe_send_embed(f"Model has been switched to {self.tagger_language.ISO_639_1}") @chatter.command(name="minutes") async def minutes(self, ctx: commands.Context, minutes: int): @@ -260,6 +329,27 @@ class Chatter(Cog): else: await ctx.send("Error occurred :(") + @chatter.command(name="trainubuntu") + async def chatter_train_ubuntu(self, ctx: commands.Context, confirmation: bool = False): + """ + WARNING: Large Download! Trains the bot using Ubuntu Dialog Corpus data. + """ + + if not confirmation: + await ctx.maybe_send_embed( + "Warning: This command downloads ~500MB then eats your CPU for training\n" + "If you're sure you want to continue, run `[p]chatter trainubuntu True`" + ) + return + + async with ctx.typing(): + future = await self.loop.run_in_executor(None, self._train_ubuntu) + + if future: + await ctx.send("Training successful!") + else: + await ctx.send("Error occurred :(") + @chatter.command(name="trainenglish") async def chatter_train_english(self, ctx: commands.Context): """ diff --git a/nudity/__init__.py b/nudity/__init__.py new file mode 100644 index 0000000..09d9dbf --- /dev/null +++ b/nudity/__init__.py @@ -0,0 +1,6 @@ +from .nudity import Nudity + + +def setup(bot): + n = Nudity(bot) + bot.add_cog(n) diff --git a/nudity/info..json b/nudity/info..json new file mode 100644 index 0000000..34c4804 --- /dev/null +++ b/nudity/info..json @@ -0,0 +1,26 @@ +{ + "author": [ + "Bobloy" + ], + "min_bot_version": [ + 3, + 3, + 11 + ], + "description": "Monitor images for NSFW content and moves them to a nsfw channel if possible", + "hidden": false, + "install_msg": "Thank you for installing Nudity. Get started with `[p]load nudity`, then `[p]help Nudity`", + "requirements": [ + "nudenet", + "tensorflow>=1.14,<2.0", + "keras>=2.4" + ], + "short": "NSFW image tracker and mover", + "tags": [ + "bobloy", + "utils", + "tools", + "nude", + "nsfw" + ] +} diff --git a/nudity/nudity.py b/nudity/nudity.py new file mode 100644 index 0000000..6eb4221 --- /dev/null +++ b/nudity/nudity.py @@ -0,0 +1,147 @@ +import pathlib + +import discord +from nudenet import NudeClassifier +from redbot.core import Config, commands +from redbot.core.bot import Red +from redbot.core.data_manager import cog_data_path + + +class Nudity(commands.Cog): + """ + V3 Cog Template + """ + + def __init__(self, bot: Red): + super().__init__() + self.bot = bot + self.config = Config.get_conf(self, identifier=9811198108111121, force_registration=True) + + default_guild = {"enabled": False, "channel_id": None} + + self.config.register_guild(**default_guild) + + # self.detector = NudeDetector() + self.classifier = NudeClassifier() + + self.data_path: pathlib.Path = cog_data_path(self) + + self.current_processes = 0 + + async def red_delete_data_for_user(self, **kwargs): + """Nothing to delete""" + return + + @commands.command(aliases=["togglenudity"], name="nudity") + async def nudity(self, ctx: commands.Context): + """Toggle nude-checking on or off""" + is_on = await self.config.guild(ctx.guild).enabled() + await self.config.guild(ctx.guild).enabled.set(not is_on) + await ctx.send("Nude checking is now set to {}".format(not is_on)) + + @commands.command() + async def nsfwchannel(self, ctx: commands.Context, channel: discord.TextChannel = None): + if channel is None: + await self.config.guild(ctx.guild).channel_id.set(None) + await ctx.send("NSFW Channel cleared") + else: + if not channel.is_nsfw(): + await ctx.send("This channel isn't NSFW!") + return + else: + await self.config.guild(ctx.guild).channel_id.set(channel.id) + await ctx.send("NSFW channel has been set to {}".format(channel.mention)) + + async def get_nsfw_channel(self, guild: discord.Guild): + channel_id = await self.config.guild(guild).channel_id() + + if channel_id is None: + return None + else: + return guild.get_channel(channel_id=channel_id) + + async def nsfw(self, message: discord.Message, images: dict): + content = message.content + guild: discord.Guild = message.guild + if not content: + content = "*`None`*" + try: + await message.delete() + except discord.Forbidden: + await message.channel.send("NSFW Image detected!") + return + + embed = discord.Embed(title="NSFW Image Detected") + embed.add_field(name="Original Message", value=content) + embed.set_author(name=message.author.name, icon_url=message.author.avatar_url) + await message.channel.send(embed=embed) + + nsfw_channel = await self.get_nsfw_channel(guild) + + if nsfw_channel is None: + return + else: + for image, r in images.items(): + if r["unsafe"] > 0.7: + await nsfw_channel.send( + "NSFW Image from {}".format(message.channel.mention), + file=discord.File(image,), + ) + + @commands.Cog.listener() + async def on_message(self, message: discord.Message): + is_private = isinstance(message.channel, discord.abc.PrivateChannel) + + if not message.attachments or is_private or message.author.bot: + # print("did not qualify") + return + + try: + is_on = await self.config.guild(message.guild).enabled() + except AttributeError: + return + + if not is_on: + print("Not on") + return + + channel: discord.TextChannel = message.channel + + if channel.is_nsfw(): + print("nsfw channel is okay") + return + + check_list = [] + for attachment in message.attachments: + # async with aiohttp.ClientSession() as session: + # img = await fetch_img(session, attachment.url) + + ext = attachment.filename + + temp_name = self.data_path / f"nudecheck{self.current_processes}_{ext}" + + self.current_processes += 1 + + print("Pre attachment save") + await attachment.save(temp_name) + check_list.append(temp_name) + + print("Pre nude check") + # nude_results = self.detector.detect(temp_name) + nude_results = self.classifier.classify([str(n) for n in check_list]) + # print(nude_results) + + if True in [r["unsafe"] > 0.7 for r in nude_results.values()]: + # print("Is nude") + await message.add_reaction("❌") + await self.nsfw(message, nude_results) + else: + # print("Is not nude") + await message.add_reaction("✅") + + +# async def fetch_img(session, url): +# with aiohttp.Timeout(10): +# async with session.get(url) as response: +# assert response.status == 200 +# return await response.read()