diff --git a/main.py b/main.py index 2a10d7c..471daef 100644 --- a/main.py +++ b/main.py @@ -7,7 +7,7 @@ from dotenv import load_dotenv from datetime import datetime, timedelta from dashboard import start_dashboard from tokenizer import Tokenizer -from model import RubyTrainer +from trainer import RubyTrainer # Load environment load_dotenv() diff --git a/model.py b/model.py index b78731d..5b99af8 100644 --- a/model.py +++ b/model.py @@ -1,9 +1,5 @@ import torch import torch.nn as nn -import torch.nn.functional as F -import os -from datetime import datetime -from collections import Counter class MiniGPT(nn.Module): @@ -26,160 +22,3 @@ class MiniGPT(nn.Module): x = block(x) x = self.ln_f(x) return self.head(x) - - -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() - self.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." - ] - - 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) - if not tokens: - return - - # Wrap with and - tokens = [self.tokenizer.vocab[""]] + tokens + [self.tokenizer.vocab[""]] - - if len(tokens) < 2: - print("[TRAIN] Skipped (not enough tokens)") - 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.2, top_k=10): - self.model.eval() - - input_ids = torch.tensor([[self.tokenizer.vocab[""]]], dtype=torch.long, device=self.device) - - token_freq = Counter() - - for _ in range(max_tokens): - with torch.no_grad(): - out = self.model(input_ids) - logits = out[:, -1, :] / temperature - - # 💡 Apply repetition penalty - for token_id, freq in token_freq.items(): - if freq > 0: - logits[0, token_id] *= 0.7 ** freq # dampens reused tokens - - probs = F.softmax(logits, dim=-1) - - if top_k > 0: - top_k_logits, top_k_indices = torch.topk(probs, top_k) - next_token = top_k_indices[0][torch.multinomial(top_k_logits, 1)] - else: - next_token = torch.multinomial(probs, 1)[0] - - token_freq[next_token.item()] += 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[""]: - break - token_ids = input_ids.squeeze(0).tolist()[1:] # skip - reply_tokens = [tid for tid in token_ids if tid != self.tokenizer.vocab.get("")] - return self.tokenizer.detokenize(reply_tokens) - - def dream(self, log_path="logs/messages.log", log_output="logs/dreams.log", max_lines=50): - print("[DREAM] Ruby is dreaming...") - - if not os.path.exists(log_path): - print("[DREAM] No memory to dream from.") - return - - with open(log_path, "r", encoding="utf-8") as f: - lines = f.readlines()[-max_lines:] - - learned = 0 - with open(log_output, "a", encoding="utf-8") as out_f: - for line in lines: - parts = line.strip().split("|") - if len(parts) >= 3: - text = parts[2].strip() - self.train_on_tokens_from_text(text) - out_f.write(f"[DREAM MEMORY] {text}\n") - learned += 1 - - print(f"[DREAM] Dream complete. Trained on {learned} memories.") - - def daydream(self, rounds=5, log_output="logs/dreams.log", say_thought=False): - print("[DAYDREAM] Ruby is imagining new thoughts...") - thoughts = [] - max_attempts = rounds * 3 # allows retries for short/empty outputs - attempts = 0 - - while len(thoughts) < rounds and attempts < max_attempts: - thought = self.generate_reply() - attempts += 1 - - if thought and len(set(thought.lower().split())) >= 3: - self.train_on_tokens_from_text(thought) - thoughts.append(thought) - - with open(log_output, "a", encoding="utf-8") as f: - for t in thoughts: - f.write(f"[DAYDREAM] {t}\n") - - # Loop dreams back into message log (optional) - 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 (in {attempts} attempts).") - - if say_thought and thoughts: - return thoughts[-1] - return None - - def reinforce_core_memory(self, log_output="logs/dreams.log"): - print("[CORE] Reinforcing Ruby's core memories...") - with open(log_output, "a", encoding="utf-8") as f: - for line in self.core_memories: - self.train_on_tokens_from_text(line) - f.write(f"[CORE MEMORY] {line}\n") diff --git a/trainer.py b/trainer.py new file mode 100644 index 0000000..a6a891c --- /dev/null +++ b/trainer.py @@ -0,0 +1,167 @@ +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[""]] + tokens + [self.tokenizer.vocab[""]] + 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[""]]], 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[""]] = 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[""]: + break + + token_ids = input_ids.squeeze(0).tolist()[1:] + reply_tokens = [t for t in token_ids if t != self.tokenizer.vocab[""]] + return self.tokenizer.detokenize(reply_tokens) + + def self_rephrase(self, original: str, max_tokens=50): + self.model.eval() + tokens = [self.tokenizer.vocab[""]] + 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[""]] = 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[""]: + break + + new_tokens = input_ids.squeeze(0).tolist()[1:] + return self.tokenizer.detokenize([t for t in new_tokens if t != self.tokenizer.vocab[""]]) + + 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)