Upload discord_bot.py
Browse files- discord_bot.py +82 -0
discord_bot.py
ADDED
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# the os module helps us access environment variables
|
2 |
+
# i.e., our API keys
|
3 |
+
import os
|
4 |
+
|
5 |
+
# these modules are for querying the Hugging Face model
|
6 |
+
import json
|
7 |
+
import requests
|
8 |
+
|
9 |
+
# the Discord Python API
|
10 |
+
import discord
|
11 |
+
|
12 |
+
# this is my Hugging Face profile link
|
13 |
+
API_URL = 'https://api-inference.huggingface.co/models/r3dhummingbird/'
|
14 |
+
|
15 |
+
class MyClient(discord.Client):
|
16 |
+
def __init__(self, model_name):
|
17 |
+
super().__init__()
|
18 |
+
self.api_endpoint = API_URL + model_name
|
19 |
+
# retrieve the secret API token from the system environment
|
20 |
+
huggingface_token = os.environ['HUGGINGFACE_TOKEN']
|
21 |
+
# format the header in our request to Hugging Face
|
22 |
+
self.request_headers = {
|
23 |
+
'Authorization': 'Bearer {}'.format(huggingface_token)
|
24 |
+
}
|
25 |
+
|
26 |
+
def query(self, payload):
|
27 |
+
"""
|
28 |
+
make request to the Hugging Face model API
|
29 |
+
"""
|
30 |
+
data = json.dumps(payload)
|
31 |
+
response = requests.request('POST',
|
32 |
+
self.api_endpoint,
|
33 |
+
headers=self.request_headers,
|
34 |
+
data=data)
|
35 |
+
ret = json.loads(response.content.decode('utf-8'))
|
36 |
+
return ret
|
37 |
+
|
38 |
+
async def on_ready(self):
|
39 |
+
# print out information when the bot wakes up
|
40 |
+
print('Logged in as')
|
41 |
+
print(self.user.name)
|
42 |
+
print(self.user.id)
|
43 |
+
print('------')
|
44 |
+
# send a request to the model without caring about the response
|
45 |
+
# just so that the model wakes up and starts loading
|
46 |
+
self.query({'inputs': {'text': 'Hello!'}})
|
47 |
+
|
48 |
+
async def on_message(self, message):
|
49 |
+
"""
|
50 |
+
this function is called whenever the bot sees a message in a channel
|
51 |
+
"""
|
52 |
+
# ignore the message if it comes from the bot itself
|
53 |
+
if message.author.id == self.user.id:
|
54 |
+
return
|
55 |
+
|
56 |
+
# form query payload with the content of the message
|
57 |
+
payload = {'inputs': {'text': message.content}}
|
58 |
+
|
59 |
+
# while the bot is waiting on a response from the model
|
60 |
+
# set the its status as typing for user-friendliness
|
61 |
+
async with message.channel.typing():
|
62 |
+
response = self.query(payload)
|
63 |
+
bot_response = response.get('generated_text', None)
|
64 |
+
|
65 |
+
# we may get ill-formed response if the model hasn't fully loaded
|
66 |
+
# or has timed out
|
67 |
+
if not bot_response:
|
68 |
+
if 'error' in response:
|
69 |
+
bot_response = '`Error: {}`'.format(response['error'])
|
70 |
+
else:
|
71 |
+
bot_response = 'Hmm... something is not right.'
|
72 |
+
|
73 |
+
# send the model's response to the Discord channel
|
74 |
+
await message.channel.send(bot_response)
|
75 |
+
|
76 |
+
def main():
|
77 |
+
# DialoGPT-medium-joshua is my model name
|
78 |
+
client = MyClient('DialoGPT-medium-joshua')
|
79 |
+
client.run(os.environ['DISCORD_TOKEN'])
|
80 |
+
|
81 |
+
if __name__ == '__main__':
|
82 |
+
main()
|