Spaces:
Runtime error
Runtime error
File size: 2,944 Bytes
3c863b8 2dd0c90 3c863b8 2dd0c90 3c863b8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
# 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['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']
# 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['MTEzNDc5NjcwMTIwNzYyOTg5NA.Ge6bnp.jHStqt72lbFrDiOV8vCzS7oDyJ0JKjuL_OdsVY'])
if __name__ == '__main__':
main() |