Mara
/

Discord_Chatbot

Model card Files Files and versions Community

Discord_Chatbot / discord_bot.py

Mara's picture

Upload discord_bot.py

6841d67 over 2 years ago

raw history blame contribute delete

No virus

2.86 kB

	# the os module helps us access environment variables
	# i.e., our API keys
	import os

	# these modules are for querying the Hugging Face model
	import json
	import requests

	# the Discord Python API
	import discord

	# this is my Hugging Face profile link
	API_URL = 'https://api-inference.huggingface.co/models/r3dhummingbird/'

	class MyClient(discord.Client):
	def __init__(self, model_name):
	super().__init__()
	self.api_endpoint = API_URL + model_name
	# retrieve the secret API token from the system environment
	huggingface_token = os.environ['HUGGINGFACE_TOKEN']
	# format the header in our request to Hugging Face
	self.request_headers = {
	'Authorization': 'Bearer {}'.format(huggingface_token)
	}

	def query(self, payload):
	"""
	make request to the Hugging Face model API
	"""
	data = json.dumps(payload)
	response = requests.request('POST',
	self.api_endpoint,
	headers=self.request_headers,
	data=data)
	ret = json.loads(response.content.decode('utf-8'))
	return ret

	async def on_ready(self):
	# print out information when the bot wakes up
	print('Logged in as')
	print(self.user.name)
	print(self.user.id)
	print('------')
	# send a request to the model without caring about the response
	# just so that the model wakes up and starts loading
	self.query({'inputs': {'text': 'Hello!'}})

	async def on_message(self, message):
	"""
	this function is called whenever the bot sees a message in a channel
	"""
	# ignore the message if it comes from the bot itself
	if message.author.id == self.user.id:
	return

	# form query payload with the content of the message
	payload = {'inputs': {'text': message.content}}

	# while the bot is waiting on a response from the model
	# set the its status as typing for user-friendliness
	async with message.channel.typing():
	response = self.query(payload)
	bot_response = response.get('generated_text', None)

	# we may get ill-formed response if the model hasn't fully loaded
	# or has timed out
	if not bot_response:
	if 'error' in response:
	bot_response = '`Error: {}`'.format(response['error'])
	else:
	bot_response = 'Hmm... something is not right.'

	# send the model's response to the Discord channel
	await message.channel.send(bot_response)

	def main():
	# DialoGPT-medium-joshua is my model name
	client = MyClient('DialoGPT-medium-joshua')
	client.run(os.environ['DISCORD_TOKEN'])

	if __name__ == '__main__':
	main()