127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
import os
|
|
import asyncio
|
|
import glob
|
|
import threading
|
|
import json
|
|
from collections import deque
|
|
import logging
|
|
import discord
|
|
|
|
from nervous_system import NervousSystem
|
|
import dashboard # <-- import your new Flask app
|
|
import brain_map # <-- import the blueprint to inject system
|
|
|
|
# Path for progress persistence
|
|
PROGRESS_PATH = 'progress.json'
|
|
|
|
# Mute logger
|
|
for noisy_logger in ["werkzeug", "flask", "flask.app"]:
|
|
logging.getLogger(noisy_logger).setLevel(logging.CRITICAL)
|
|
|
|
# ─── Initialize Ruby & Discord ───────────────────────────────────────────────── # noqa: E501
|
|
|
|
intents = discord.Intents.default()
|
|
intents.message_content = True
|
|
client = discord.Client(intents=intents)
|
|
system = NervousSystem()
|
|
system.history = deque(maxlen=100)
|
|
|
|
# Load or resume vocab + embeddings
|
|
system.sensory.load_vocab('vocab.json')
|
|
system._resize_embeddings()
|
|
print('Loaded vocab size:', len(system.sensory.stoi))
|
|
|
|
# Resume progress
|
|
if os.path.isfile(PROGRESS_PATH):
|
|
with open(PROGRESS_PATH, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
system.processed_lines = data.get('processed_lines', 0)
|
|
else:
|
|
system.processed_lines = 0
|
|
|
|
# Compute total book lines
|
|
total = sum(
|
|
1
|
|
for path in glob.glob('books/*.txt')
|
|
for line in open(path, encoding='utf-8')
|
|
if line.strip()
|
|
)
|
|
system.total_lines = total
|
|
|
|
print(f'Resuming training at {system.processed_lines}/{system.total_lines} lines')
|
|
|
|
# Inject into Flask contexts
|
|
dashboard.system = system
|
|
brain_map.system = system
|
|
|
|
# ─── Book-training when idle ──────────────────────────────────────────────────── # noqa: E501
|
|
|
|
|
|
async def train_books_idle():
|
|
await client.wait_until_ready()
|
|
await asyncio.sleep(5)
|
|
processed = 0
|
|
skip = system.processed_lines
|
|
|
|
for path in glob.glob('books/*.txt'):
|
|
with open(path, encoding='utf-8') as f:
|
|
for raw in f:
|
|
text = raw.strip()
|
|
if not text:
|
|
continue
|
|
if processed < skip:
|
|
processed += 1
|
|
continue
|
|
|
|
await asyncio.to_thread(system.train, text, text)
|
|
processed += 1
|
|
system.processed_lines = processed
|
|
|
|
if processed % 200 == 0 or processed == system.total_lines:
|
|
system.sensory.save_vocab('vocab.json')
|
|
with open(PROGRESS_PATH, 'w', encoding='utf-8') as pf:
|
|
json.dump({'processed_lines': processed}, pf)
|
|
|
|
# Final checkpoint
|
|
system.sensory.save_vocab('vocab.json')
|
|
with open(PROGRESS_PATH, 'w', encoding='utf-8') as pf:
|
|
json.dump({'processed_lines': system.processed_lines}, pf)
|
|
|
|
|
|
@client.event
|
|
async def on_ready():
|
|
print(f'Ruby is online as {client.user}!')
|
|
asyncio.create_task(train_books_idle())
|
|
|
|
|
|
@client.event
|
|
async def on_message(message: discord.Message):
|
|
if message.author == client.user or not message.content:
|
|
return
|
|
|
|
user_text = message.content.strip()
|
|
reply = system.generate(user_text)
|
|
await message.channel.send(reply)
|
|
|
|
system.history.append({'user': user_text, 'bot': reply})
|
|
asyncio.create_task(asyncio.to_thread(system.train, user_text, reply))
|
|
|
|
|
|
# ─── Launch Dashboard & Bot ──────────────────────────────────────────────────── # noqa: E501
|
|
|
|
|
|
def run_dashboard():
|
|
dashboard.app.run(
|
|
host='0.0.0.0', port=5000,
|
|
debug=False, use_reloader=False
|
|
)
|
|
|
|
|
|
threading.Thread(target=run_dashboard, daemon=True).start()
|
|
print('Dashboard available at http://127.0.0.1:5000')
|
|
|
|
token = os.getenv('DISCORD_TOKEN')
|
|
if not token:
|
|
raise RuntimeError('Please set the DISCORD_TOKEN environment variable')
|
|
client.run(token)
|