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 info
Just 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 online
Alpha release, please report bugs |
-| nudity | **Incomplete** | Checks for NSFW images posted in non-NSFW channels
Library this is based on has a bug, waiting for author to merge my PR |
+| nudity | **Alpha** | Checks for NSFW images posted in non-NSFW channels
Switched 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 server
Alpha release, please report any bugs |
| reactrestrict | **Alpha** | Removes reactions by role per channel
A 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()