Ruby/model/rehearsal.py

31 lines
834 B
Python

import torch
from model.brain import model, tokenizer, DEVICE
from model.trainer import train_on_message
from model.dynamic_expand import expand_model_if_needed
def simulate_conversation():
expand_model_if_needed()
model.eval()
max_token_id = model.head.out_features - 1
if max_token_id < 1:
return # Safeguard if model is still too small
seed = torch.randint(0, max_token_id + 1, (1, 5), device=DEVICE)
seed = seed[:, -128:] # Clamp sequence length
output = model(seed)
preds = torch.argmax(output, dim=-1).squeeze().tolist()
if isinstance(preds, int):
preds = [preds]
# 🛡 Clamp predictions too
preds = [min(max(p, 0), max_token_id) for p in preds]
text = tokenizer.detokenize(preds)
if text and len(text.split()) >= 3:
train_on_message(text)