65 lines
1.6 KiB
Python
65 lines
1.6 KiB
Python
import os
|
|
import discord
|
|
import torch
|
|
from dotenv import load_dotenv
|
|
from train_gpt_model import process_message
|
|
from gpt_model import load_model
|
|
|
|
# Load environment variables from .env file
|
|
load_dotenv()
|
|
|
|
# Get the Discord bot token from environment variables
|
|
TOKEN = os.getenv("DISCORD_TOKEN")
|
|
|
|
# Load the vocabulary
|
|
with open("vocab.txt", "r", encoding="utf-8") as f:
|
|
text = f.read()
|
|
chars = sorted(list(set(text)))
|
|
|
|
# Ensure that space and other special characters are included
|
|
required_chars = " \n\r\t"
|
|
for char in required_chars:
|
|
if char not in chars:
|
|
chars.append(char)
|
|
|
|
# Add a special token for unknown characters
|
|
special_token = "<unk>"
|
|
if special_token not in chars:
|
|
chars.append(special_token)
|
|
|
|
vocab_size = len(chars)
|
|
string_to_int = {ch: i for i, ch in enumerate(chars)}
|
|
int_to_string = {i: ch for i, ch in enumerate(chars)}
|
|
|
|
# Initialize and load the model
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model = load_model(vocab_size, "phoebe_model.pt").to(device)
|
|
|
|
# Initialize Discord client
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
client = discord.Client(intents=intents)
|
|
|
|
|
|
@client.event
|
|
async def on_ready():
|
|
print(f"We have logged in as {client.user}")
|
|
|
|
|
|
@client.event
|
|
async def on_message(message):
|
|
if message.author == client.user:
|
|
return
|
|
|
|
# Debug: print the message content
|
|
print(f"Received message: '{message.content}'")
|
|
|
|
# Process the message and get a response
|
|
response = process_message(message.content)
|
|
|
|
# Send the response back to the Discord channel
|
|
await message.channel.send(response)
|
|
|
|
|
|
client.run(TOKEN)
|