290 lines
9.6 KiB
Python
290 lines
9.6 KiB
Python
import asyncio
|
|
import discord
|
|
from discord.ext import commands
|
|
import numpy as np
|
|
import random
|
|
|
|
class Chatbot(object):
|
|
"""Een heuse chatbot. Under construction..."""
|
|
|
|
def __init__(self, bot):
|
|
self.bot = bot
|
|
self.voice_states = {}
|
|
self.END_OF_MESSAGE = "<EOM>"
|
|
self.START_OF_MESSAGE = "<SOM>"
|
|
try:
|
|
np.load('markov_dict.npy')
|
|
except FileNotFoundError:
|
|
#If there's no initial file, make a new one:
|
|
initFile = {'pairs':{self.START_OF_MESSAGE:{'Hoi':1},'Hoi':{self.END_OF_MESSAGE:1}}, 'trairs':{}, 'responsePairs':{}}
|
|
np.save("markov_dict.npy", initFile)
|
|
except:
|
|
print("ERROR! - Could not open or create Markov chat file")
|
|
|
|
|
|
#Do stuff when a message comes in:
|
|
@self.bot.event
|
|
@asyncio.coroutine
|
|
def on_message(message):
|
|
if message.author.bot:
|
|
return
|
|
|
|
if message.content.strip()[0] != '!': #don't respond to a command with a chat
|
|
if message.content.lower() == "nihao":
|
|
yield from message.channel.send("kankerlauw")
|
|
return
|
|
|
|
if "DMChannel" in type(message.channel).__name__: #respond to messages in PM channels
|
|
yield from self.respondToMessage(message)
|
|
else:
|
|
pass #TODO respond to other channels at random
|
|
|
|
#process commands:
|
|
yield from self.bot.process_commands(message)
|
|
|
|
def __check(self, ctx):
|
|
return True
|
|
|
|
|
|
|
|
#Make word pairs:
|
|
def makePairs(self, corpus):
|
|
for i in range(len(corpus)-1):
|
|
yield (corpus[i], corpus[i+1])
|
|
|
|
#Make word trairs. Yes, I completely made up "trair"
|
|
#It's like pairs, but with three, but not really, because there are two. OK?
|
|
#[1 0 1]
|
|
def makeTrairs(self, corpus):
|
|
for i in range(1, len(corpus)-1):
|
|
yield (corpus[i-1], corpus[i+1])
|
|
|
|
#Learn from a human message
|
|
#Returns an updated dict
|
|
def learnFrom(self, wordDictPairs, wordDictTrairs, message: str):
|
|
newPairDict = wordDictPairs
|
|
newTrairDict = wordDictTrairs
|
|
|
|
#If the message is only a tag, don't learn from it. return
|
|
if message.count(' ') == 0 and message[0:2] == "<@":
|
|
return newPairDict, newTrairDict;
|
|
|
|
message = self.START_OF_MESSAGE + " " + message + " " + self.END_OF_MESSAGE
|
|
corpus = message.split()
|
|
for n, word in enumerate(corpus):
|
|
if word == "@everyone":
|
|
corpus[n] = "@iedereen"
|
|
pairs = self.makePairs(corpus)
|
|
trairs = self.makeTrairs(corpus)
|
|
|
|
#Add wordcounts to pair dict
|
|
for word_1, word_2 in pairs:
|
|
if word_1 in newPairDict.keys():
|
|
if word_2 in newPairDict[word_1].keys():
|
|
newPairDict[word_1][word_2] += 1
|
|
else:
|
|
newPairDict[word_1][word_2] = 1
|
|
else:
|
|
newPairDict[word_1] = {word_2 : 1}
|
|
#Add wordcounts to trair dict
|
|
for word_1, word_3 in trairs:
|
|
if word_1 in newTrairDict.keys():
|
|
if word_3 in newTrairDict[word_1].keys():
|
|
newTrairDict[word_1][word_3] += 1
|
|
else:
|
|
newTrairDict[word_1][word_3] = 1
|
|
else:
|
|
newTrairDict[word_1] = {word_3 : 1}
|
|
|
|
return newPairDict, newTrairDict;
|
|
|
|
#Learn to respond to messages
|
|
#Returns ???
|
|
def learnResponse(self, wordDictResp, message1: str, message2: str):
|
|
STOP_WORDS = ["aan", "achter", "al", "ben", "dan", "dat", "de", "die", "dit", "een", "en", "er", "gehad", "had", "heb",
|
|
"hen", "het", "hun", "in", "is", "kon", "met", "na", "of", "om", "ook", "op", "tot", "van", "was", "wat", "zo",
|
|
"a", "an", "the", "am", "was", "were", "be", "not", "at"]
|
|
|
|
newRespPairDict = wordDictResp
|
|
|
|
#Split messages into words and remove stopwords:
|
|
message1List = []
|
|
message2List = []
|
|
for word in message1.split():
|
|
if word not in STOP_WORDS:
|
|
message1List.append(word)
|
|
for word in message2.split():
|
|
if word not in STOP_WORDS:
|
|
message2List.append(word)
|
|
|
|
#link words of response to words of initial message
|
|
for word_1 in message1List:
|
|
for word_2 in message2List:
|
|
if word_1 in newRespPairDict.keys():
|
|
if word_2 in newRespPairDict[word_1].keys():
|
|
newRespPairDict[word_1][word_2] += 1
|
|
else:
|
|
newRespPairDict[word_1][word_2] = 1
|
|
else:
|
|
newRespPairDict[word_1] = {word_2 : 1}
|
|
|
|
return newRespPairDict
|
|
|
|
|
|
|
|
@commands.command(pass_context=True, hidden=True)
|
|
@asyncio.coroutine
|
|
def trainmarkovchat(self, ctx, channelId=409751773595566081, nrOfMessages=50):
|
|
"""Train that chain, jermaine
|
|
param channelId: use this channel to train
|
|
param nrOfMessages: use the x most recent messages
|
|
"""
|
|
|
|
channel=self.bot.get_channel(channelId)
|
|
#Open current chain:
|
|
wordDict = np.load('markov_dict.npy').item()
|
|
|
|
#Iterate through the last x messages from the channel:
|
|
messageIterator = channel.history(limit=nrOfMessages, before=None, after=None, reverse=False, around=None)
|
|
allMessages = []
|
|
for i in range(nrOfMessages):
|
|
msg = yield from messageIterator.next()
|
|
allMessages.append(msg)
|
|
|
|
i = 0
|
|
nextMessageIndex = -10
|
|
#Note: the loop goes from new to old
|
|
for msg in allMessages:
|
|
messageContent = msg.content.replace('(', '').replace(')', '').replace('"', '')
|
|
messageContent = messageContent.strip()
|
|
|
|
if messageContent == "":
|
|
pass # Bericht is leeg. Negeer.
|
|
elif msg.author.bot:
|
|
pass # Dit is een bot chat. Negeer.
|
|
elif messageContent[0] == "!":
|
|
pass # Dit is een commando. Negeer.
|
|
elif "```" in messageContent:
|
|
pass # Daar zit een codeblok in. Negeer.
|
|
elif messageContent[0:7] == "http://" or messageContent[0:8] == "https://":
|
|
pass # Een link. Negeer. TODO: post af en toe een random link die geen repost is
|
|
else:
|
|
# Usable message!
|
|
wordDict['pairs'], wordDict['trairs'] = self.learnFrom(wordDict['pairs'], wordDict['trairs'], messageContent)
|
|
|
|
if (i-nextMessageIndex) == 1 and msg.author != nextMessage.author: #the previous usable messages was the one right after this one and was not from the same author
|
|
timeDelta = nextMessage.created_at - msg.created_at
|
|
#TODO check for different author
|
|
if timeDelta.days == 0 and timeDelta.seconds < 3600: #less than an hour difference
|
|
#this is counted as a response. learn from it
|
|
wordDict['responsePairs'] = self.learnResponse(wordDict['responsePairs'], messageContent, nextMessageContent)
|
|
|
|
#save for next iteration:
|
|
nextMessage = msg
|
|
nextMessageContent = messageContent
|
|
nextMessageIndex = i
|
|
i += 1
|
|
|
|
np.save("markov_dict.npy", wordDict)
|
|
|
|
yield from ctx.channel.send("Ik heb weer wat nieuwe woordjes geleerd. :slight_smile:")
|
|
|
|
|
|
@commands.command(pass_context=True, hidden=False)
|
|
@asyncio.coroutine
|
|
def reageer(self, ctx):
|
|
"""UNDER CONSTRUCTION. Laat Stroop ergens op reageren.
|
|
Stroop leert van onze gesprekken. Hij kan zich nu voordoen als een van ons. Soort van.
|
|
"""
|
|
|
|
#Get previous message to respond to
|
|
messageIterator = ctx.channel.history(limit=3, before=None, after=None, reverse=False, around=None)
|
|
messageToRespondTo = yield from messageIterator.next() #this will be the !reageer command. skip
|
|
messageToRespondTo = yield from messageIterator.next()
|
|
if messageToRespondTo.content.strip()[0] == '!' or messageToRespondTo.author.bot: #another chance
|
|
messageToRespondTo = yield from messageIterator.next() #but it's not a huge problem if he responds to himself
|
|
if messageToRespondTo.content.strip()[0] == '!':
|
|
wordsToRespondTo = []
|
|
else:
|
|
wordsToRespondTo = messageToRespondTo.content.strip().split()
|
|
#print("respond to: ", wordsToRespondTo)
|
|
|
|
#Delete command
|
|
try:
|
|
yield from ctx.message.delete()
|
|
except:
|
|
pass
|
|
|
|
response = self.getResponseForMessage(wordsToRespondTo)
|
|
yield from ctx.channel.send(response)
|
|
|
|
#Respond to a message.
|
|
#Assumes the message isn't a command
|
|
def respondToMessage(self, message):
|
|
wordsToRespondTo = message.content.strip().split()
|
|
|
|
response = self.getResponseForMessage(wordsToRespondTo)
|
|
yield from message.channel.send(response)
|
|
|
|
#Gets a response
|
|
def getResponseForMessage(self, wordsToRespondTo):
|
|
|
|
MAX_N_WORDS = 100 #Don't use more words than this
|
|
|
|
#Get dictionary/Markov chain
|
|
wordDict = np.load('markov_dict.npy').item()
|
|
wordPairDict = wordDict['pairs']
|
|
wordTrairDict = wordDict['trairs']
|
|
responsePairDict = wordDict['responsePairs']
|
|
#TODO store in global variable to minimise IO?
|
|
|
|
#Wtart with <SOM>
|
|
chain = [self.START_OF_MESSAGE]
|
|
|
|
#Go through the chain:
|
|
for i in range(1, MAX_N_WORDS):
|
|
if chain[-1] == self.END_OF_MESSAGE:
|
|
break
|
|
|
|
#Find the possibilities:
|
|
pairPossibilities = wordPairDict[chain[-1]].copy()
|
|
|
|
#Increase score if it would be a good response:
|
|
for userWord in wordsToRespondTo:
|
|
if userWord in responsePairDict:
|
|
for goodResponseWord in responsePairDict[userWord]:
|
|
if goodResponseWord in pairPossibilities.keys():
|
|
pairPossibilities[goodResponseWord] += responsePairDict[userWord][goodResponseWord]
|
|
#Increase score for trairs:
|
|
if i > 1: #we can only compare trairs if we have at least two words already
|
|
if chain[-2] in wordTrairDict:
|
|
trairPossibilities = wordTrairDict[chain[-2]].copy()
|
|
#If a trair matches, increase the frequency of the pair
|
|
for word2 in pairPossibilities.keys():
|
|
if word2 in trairPossibilities.keys():
|
|
pairPossibilities[word2] += trairPossibilities[word2]
|
|
|
|
#Calculate probabilities and pick a next word:
|
|
#https://stackoverflow.com/questions/835092/python-dictionary-are-keys-and-values-always-the-same-order
|
|
possible_words = list(pairPossibilities.keys())
|
|
weights = np.array(list(pairPossibilities.values()))
|
|
probs = weights / weights.sum()
|
|
chain.append(np.random.choice(possible_words, p=probs))
|
|
|
|
|
|
if chain[0] == self.START_OF_MESSAGE:
|
|
chain = chain[1:-1] #remove <SOM> and <EOM>
|
|
else:
|
|
chain = chain[:-1] #remove <EOM>
|
|
|
|
#remove 90% of tags:
|
|
for n, word in enumerate(chain):
|
|
if word[0:2] == "<@":
|
|
a = random.randint(1,10)
|
|
if a != 1:
|
|
chain[n] = "@jemoeder"
|
|
|
|
#print(' '.join(chain))
|
|
#yield from ctx.channel.send(' '.join(chain))
|
|
return ' '.join(chain)
|
|
|