Switched over to using a seperate train code instead of it being in the model code.

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()

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 <START> and <END>
-        tokens = [self.tokenizer.vocab["<START>"]] + tokens + [self.tokenizer.vocab["<END>"]]
-
-        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["<START>"]]], 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["<END>"]:
-                    break
-        token_ids = input_ids.squeeze(0).tolist()[1:]  # skip <START>
-        reply_tokens = [tid for tid in token_ids if tid != self.tokenizer.vocab.get("<END>")]
-        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")

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["<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)