81 lines
2.6 KiB
Python
81 lines
2.6 KiB
Python
import discord
|
|
import asyncio
|
|
import os
|
|
from dotenv import load_dotenv
|
|
from datetime import datetime, timedelta
|
|
from tokenizer import Tokenizer
|
|
from model import RubyTrainer
|
|
|
|
# Load environment
|
|
load_dotenv()
|
|
TOKEN = os.getenv("DISCORD_TOKEN")
|
|
|
|
if not TOKEN:
|
|
raise RuntimeError("Bot token not found in .env")
|
|
|
|
# Setup intents
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
intents.dm_messages = True
|
|
intents = intents
|
|
|
|
|
|
class Ruby(discord.Client):
|
|
def __init__(self):
|
|
super().__init__(intents=intents)
|
|
self.tokenizer = Tokenizer()
|
|
self.trainer = RubyTrainer(self.tokenizer)
|
|
self.log_path = os.path.join("logs", "messages.log")
|
|
self.last_message_time = datetime.utcnow()
|
|
self.idle_threshold = timedelta(seconds=120) # adjust as needed
|
|
self.loop.create_task(self.idle_dream_loop())
|
|
os.makedirs("logs", exist_ok=True)
|
|
|
|
async def on_ready(self):
|
|
print(f"[READY] Logged in as {self.user} (ID: {self.user.id})")
|
|
|
|
async def idle_dream_loop(self):
|
|
await self.wait_until_ready()
|
|
while not self.is_closed():
|
|
now = datetime.utcnow()
|
|
if now - self.last_message_time > self.idle_threshold:
|
|
print("[IDLE] Ruby has been idle — entering dream mode.")
|
|
self.trainer.dream()
|
|
self.trainer.daydream()
|
|
self.last_message_time = datetime.utcnow() # reset after dreaming
|
|
await asyncio.sleep(30) # check every 30 seconds
|
|
|
|
async def on_message(self, message: discord.Message):
|
|
if message.author.id == self.user.id:
|
|
return
|
|
|
|
self.log_message(message)
|
|
self.trainer.train_on_tokens_from_text(message.content.strip())
|
|
|
|
reply = self.trainer.generate_reply()
|
|
if reply.strip():
|
|
await message.channel.send(reply)
|
|
else:
|
|
print("[REPLY] Skipped (empty)")
|
|
|
|
|
|
def log_message(self, message: discord.Message):
|
|
timestamp = datetime.utcnow().isoformat()
|
|
log_entry = f"{timestamp} | {message.author.name} | {message.content.strip()}\n"
|
|
|
|
with open(self.log_path, "a", encoding="utf-8") as f:
|
|
f.write(log_entry)
|
|
|
|
print(f"[LOGGED] {log_entry.strip()}")
|
|
|
|
def train_on_message(self, message: discord.Message):
|
|
text = message.content.strip()
|
|
self.trainer.train_on_tokens_from_text(text)
|
|
token_tensor = torch.tensor(tokens, dtype=torch.long)
|
|
loss = train_on_tokens(self.model, tokens, self.optimizer, self.criterion, device="cpu")
|
|
print(f"[TRAIN] Tokens: {tokens} | Loss: {loss:.4f}")
|
|
|
|
# Run Ruby
|
|
client = Ruby()
|
|
client.run(TOKEN)
|