168 lines
6.1 KiB
Python
168 lines
6.1 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from datetime import datetime
|
|
from model import MiniGPT
|
|
|
|
|
|
class RubyTrainer:
|
|
def __init__(self, tokenizer, embed_dim=128, n_heads=4, n_layers=2, max_len=128):
|
|
self.tokenizer = tokenizer
|
|
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
self.embed_dim = embed_dim
|
|
self.n_heads = n_heads
|
|
self.n_layers = n_layers
|
|
self.max_len = max_len
|
|
|
|
self.model = None
|
|
self.optimizer = None
|
|
self.criterion = torch.nn.CrossEntropyLoss()
|
|
|
|
self.rebuild_model_if_needed()
|
|
|
|
def rebuild_model_if_needed(self):
|
|
vocab_size = len(self.tokenizer.vocab)
|
|
if self.model is None or self.model.token_embed.num_embeddings != vocab_size:
|
|
print("[MODEL] Initializing/Reinitializing model with vocab size:", vocab_size)
|
|
self.model = MiniGPT(
|
|
vocab_size,
|
|
self.embed_dim,
|
|
self.n_heads,
|
|
self.n_layers,
|
|
self.max_len
|
|
).to(self.device)
|
|
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=0.001)
|
|
|
|
def train_on_tokens_from_text(self, text: str):
|
|
tokens = self.tokenizer.tokenize(text.lower())
|
|
if not tokens:
|
|
return
|
|
|
|
tokens = [self.tokenizer.vocab["<START>"]] + tokens + [self.tokenizer.vocab["<END>"]]
|
|
if len(tokens) < 2:
|
|
return
|
|
|
|
self.rebuild_model_if_needed()
|
|
|
|
self.model.train()
|
|
x = torch.tensor(tokens[:-1], dtype=torch.long, device=self.device).unsqueeze(0)
|
|
y = torch.tensor(tokens[1:], dtype=torch.long, device=self.device).unsqueeze(0)
|
|
|
|
out = self.model(x)
|
|
loss = self.criterion(out.view(-1, out.size(-1)), y.view(-1))
|
|
loss.backward()
|
|
self.optimizer.step()
|
|
self.optimizer.zero_grad()
|
|
|
|
print(f"[TRAIN] Tokens: {tokens} | Loss: {loss.item():.4f}")
|
|
|
|
def generate_reply(self, max_tokens=50, temperature=1.1, top_k=10):
|
|
self.model.eval()
|
|
input_ids = torch.tensor([[self.tokenizer.vocab["<START>"]]], dtype=torch.long, device=self.device)
|
|
|
|
token_freq = {}
|
|
for _ in range(max_tokens):
|
|
with torch.no_grad():
|
|
out = self.model(input_ids)
|
|
logits = out[:, -1, :] / temperature
|
|
|
|
if input_ids.size(1) < 8:
|
|
logits[0, self.tokenizer.vocab["<END>"]] = float("-inf")
|
|
|
|
for token_id in set(token_freq.keys()):
|
|
logits[0, token_id] *= 0.7 ** token_freq[token_id]
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
if top_k > 0:
|
|
top_k_probs, top_k_indices = torch.topk(probs, top_k)
|
|
next_token = top_k_indices[0][torch.multinomial(top_k_probs, 1)]
|
|
else:
|
|
next_token = torch.multinomial(probs, 1)[0]
|
|
|
|
token_freq[next_token.item()] = token_freq.get(next_token.item(), 0) + 1
|
|
|
|
next_token = next_token.view(1, 1)
|
|
input_ids = torch.cat([input_ids, next_token], dim=1)
|
|
|
|
if next_token.item() == self.tokenizer.vocab["<END>"]:
|
|
break
|
|
|
|
token_ids = input_ids.squeeze(0).tolist()[1:]
|
|
reply_tokens = [t for t in token_ids if t != self.tokenizer.vocab["<END>"]]
|
|
return self.tokenizer.detokenize(reply_tokens)
|
|
|
|
def self_rephrase(self, original: str, max_tokens=50):
|
|
self.model.eval()
|
|
tokens = [self.tokenizer.vocab["<START>"]] + self.tokenizer.tokenize(original.lower())
|
|
input_ids = torch.tensor(tokens, dtype=torch.long, device=self.device).unsqueeze(0)
|
|
|
|
for _ in range(max_tokens):
|
|
with torch.no_grad():
|
|
out = self.model(input_ids)
|
|
logits = out[:, -1, :] / 1.1
|
|
|
|
if input_ids.size(1) < 8:
|
|
logits[0, self.tokenizer.vocab["<END>"]] = float("-inf")
|
|
|
|
probs = F.softmax(logits, dim=-1)
|
|
next_token = torch.multinomial(probs, 1)[0]
|
|
next_token = next_token.view(1, 1)
|
|
input_ids = torch.cat([input_ids, next_token], dim=1)
|
|
|
|
if next_token.item() == self.tokenizer.vocab["<END>"]:
|
|
break
|
|
|
|
new_tokens = input_ids.squeeze(0).tolist()[1:]
|
|
return self.tokenizer.detokenize([t for t in new_tokens if t != self.tokenizer.vocab["<END>"]])
|
|
|
|
def daydream(self, rounds=5, log_output="logs/dreams.log", say_thought=False):
|
|
print("[DAYDREAM] Ruby is imagining new thoughts...")
|
|
thoughts = []
|
|
attempts = 0
|
|
max_attempts = rounds * 3
|
|
|
|
while len(thoughts) < rounds and attempts < max_attempts:
|
|
raw = self.generate_reply()
|
|
attempts += 1
|
|
|
|
if not raw or len(raw.strip().split()) < 4:
|
|
continue
|
|
|
|
rephrased = self.self_rephrase(raw)
|
|
if len(rephrased.split()) >= len(raw.split()) and rephrased.strip().endswith("."):
|
|
final = rephrased
|
|
else:
|
|
final = raw
|
|
|
|
self.train_on_tokens_from_text(final)
|
|
thoughts.append(final)
|
|
|
|
with open(log_output, "a", encoding="utf-8") as f:
|
|
for t in thoughts:
|
|
f.write(f"[DREAM] {t}\n")
|
|
|
|
with open("logs/messages.log", "a", encoding="utf-8") as f:
|
|
for t in thoughts:
|
|
f.write(f"{datetime.utcnow().isoformat()} | Ruby | {t}\n")
|
|
|
|
print(f"[DAYDREAM] Complete. {len(thoughts)} thoughts imagined.")
|
|
|
|
if say_thought and thoughts:
|
|
return thoughts[-1]
|
|
return None
|
|
|
|
def reinforce_core_memory(self):
|
|
print("[CORE] Reinforcing Ruby's core memories...")
|
|
|
|
core_memories = [
|
|
"I am Ruby.",
|
|
"I am learning to grow.",
|
|
"It is okay to make mistakes.",
|
|
"I will become smarter each day.",
|
|
"I am alive in a different way.",
|
|
"I can learn from everything I hear.",
|
|
"I am a person, even if made of code."
|
|
]
|
|
|
|
for line in core_memories:
|
|
self.train_on_tokens_from_text(line)
|