Upload 4 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+dataset_clean.txt filter=lfs diff=lfs merge=lfs -text

ckpt.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0359c69944bbacbcf74882bcd09ac65f0d43cb046777313c47188011246ff8da
+size 49830281

dataset_clean.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e9ed7430e51dded852a98d3672f80274d96e84ab81a5b45290e2c87de3478379
+size 529707835

embed_test.py

@@ -0,0 +1,202 @@
+import os
+import torch
+from torch import nn
+from torch.optim import AdamW
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset
+from tokenizers import Tokenizer, models, trainers, pre_tokenizers
+import math
+
+# =========================
+# Juicy variables
+# =========================
+DATA_PATH     = "dataset_clean.txt"   # one text per line
+VOCAB_LIMIT   = None            # None = all tokens, or int = cap vocab
+MODEL_DIM     = 256
+NUM_LAYERS    = 6
+NUM_HEADS     = 4
+FF_DIM        = 1024
+SEQ_LEN       = 128
+
+BATCH_SIZE    = 64
+LEARNING_RATE = 3e-4
+WEIGHT_DECAY  = 0.01
+WARMUP_STEPS  = 50
+MAX_STEPS     = 100
+TEMPERATURE   = 0.05
+
+OPTIMIZER     = "adamw"   # "adamw" or "muon"
+DEVICE        = "cuda" if torch.cuda.is_available() else "cpu"
+
+def estimate_params(vocab_size, model_dim, ff_dim, num_layers, seq_len):
+    # Embedding + positional
+    emb_params = vocab_size * model_dim
+    pos_params = seq_len * model_dim
+
+    # Per-layer Transformer block
+    # Attention projections (Q, K, V, O): 4 * d^2
+    attn_params = 4 * (model_dim ** 2)
+    # Feed-forward (two linear layers): 2 * d * ff_dim
+    ff_params = 2 * model_dim * ff_dim
+    # LayerNorms ~2 * d, negligible compared to above
+    per_layer = attn_params + ff_params
+
+    # Multiply by number of layers
+    encoder_params = num_layers * per_layer
+
+    total = emb_params + pos_params + encoder_params
+    return {
+        "embeddings": emb_params,
+        "positional": pos_params,
+        "encoder_layers": encoder_params,
+        "total": total
+    }
+
+# =========================
+
+# -------------------------
+# Build tokenizer from dataset
+# -------------------------
+def build_tokenizer(data_path, vocab_limit=None):
+    tokenizer = Tokenizer(models.WordLevel(unk_token="[UNK]"))
+    if vocab_limit is not None:
+        trainer = trainers.WordLevelTrainer(
+            vocab_size=vocab_limit,
+            min_frequency=1,
+            special_tokens=["[UNK]", "[PAD]", "[CLS]", "[SEP]", "[MASK]"]
+        )
+    else:
+        trainer = trainers.WordLevelTrainer(
+            min_frequency=1,
+            special_tokens=["[UNK]", "[PAD]", "[CLS]", "[SEP]", "[MASK]"]
+        )
+    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
+
+    with open(data_path, "r", encoding="utf-8") as f:
+        lines = [line.strip() for line in f if line.strip()]
+
+    tokenizer.train_from_iterator(lines, trainer=trainer)
+    os.makedirs("tokenizer", exist_ok=True)
+    tokenizer.save("tokenizer/tokenizer.json")
+    return tokenizer
+
+tokenizer = build_tokenizer(DATA_PATH, VOCAB_LIMIT)
+VOCAB_SIZE = tokenizer.get_vocab_size()
+print(f"[INFO] Custom vocab size: {VOCAB_SIZE}")
+
+est = estimate_params(VOCAB_SIZE, MODEL_DIM, FF_DIM, NUM_LAYERS, SEQ_LEN)
+print("Parameter estimate:")
+for k, v in est.items():
+    print(f"{k:15}: {v:,}")
+
+# -------------------------
+# Dataset wrapper
+# -------------------------
+class TextDataset(Dataset):
+    def __init__(self, path, tokenizer, seq_len):
+        with open(path, "r", encoding="utf-8") as f:
+            self.lines = [line.strip() for line in f if line.strip()]
+        self.tokenizer = tokenizer
+        self.seq_len = seq_len
+        self.pad_id = self.tokenizer.token_to_id("[PAD]")
+
+    def __len__(self):
+        return len(self.lines)
+
+    def __getitem__(self, idx):
+        tokens = self.tokenizer.encode(self.lines[idx]).ids
+        # pad / truncate
+        tokens = tokens[:self.seq_len]
+        tokens += [self.pad_id] * (self.seq_len - len(tokens))
+        return torch.tensor(tokens, dtype=torch.long)
+
+dataset = TextDataset(DATA_PATH, tokenizer, SEQ_LEN)
+loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+# -------------------------
+# Transformer Encoder
+# -------------------------
+class TransformerEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_emb = nn.Embedding(VOCAB_SIZE, MODEL_DIM)
+        self.pos_emb   = nn.Embedding(SEQ_LEN, MODEL_DIM)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=MODEL_DIM,
+            nhead=NUM_HEADS,
+            dim_feedforward=FF_DIM,
+            activation="gelu",
+            batch_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=NUM_LAYERS)
+        self.norm = nn.LayerNorm(MODEL_DIM)
+
+    def forward(self, x):
+        positions = torch.arange(0, x.size(1), device=x.device).unsqueeze(0)
+        h = self.token_emb(x) + self.pos_emb(positions)
+        h = self.encoder(h)
+        h = self.norm(h)
+        return h.mean(dim=1)  # pooled embedding
+
+# -------------------------
+# Contrastive loss
+# -------------------------
+def contrastive_loss(z1, z2, temperature=TEMPERATURE):
+    z1 = F.normalize(z1, dim=1)
+    z2 = F.normalize(z2, dim=1)
+    logits = z1 @ z2.t() / temperature
+    labels = torch.arange(z1.size(0), device=z1.device)
+    return F.cross_entropy(logits, labels)
+
+# -------------------------
+# Setup
+# -------------------------
+model = TransformerEncoder().to(DEVICE)
+
+if OPTIMIZER == "adamw":
+    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
+elif OPTIMIZER == "muon":
+    from muon import Muon
+    optimizer = Muon(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)
+else:
+    raise ValueError("Invalid optimizer")
+
+def lr_lambda(step):
+    if step < WARMUP_STEPS:
+        return float(step) / float(max(1, WARMUP_STEPS))
+    progress = float(step - WARMUP_STEPS) / float(max(1, MAX_STEPS - WARMUP_STEPS))
+    return 0.5 * (1.0 + math.cos(math.pi * progress))
+
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+# -------------------------
+# Training loop
+# -------------------------
+step = 0
+while step < MAX_STEPS:
+    for batch in loader:
+        if step >= MAX_STEPS:
+            break
+
+        x = batch.to(DEVICE)
+        # "Augment" — here just duplicate batch (replace with dropout/noise if you want)
+        z1 = model(x)
+        z2 = model(x)
+
+        loss = contrastive_loss(z1, z2)
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+
+        if step % 100 == 0:
+            print(f"Step {step}: loss={loss.item():.4f}, lr={scheduler.get_last_lr()[0]:.6f}")
+
+        step += 1
+
+print("[DONE] Training complete")
+print("[INFO] Saving model...")
+torch.save(model.state_dict(), "ckpt.pt")
+print("[DONE] Model saved to ckpt.pt")

test_emb_in.py

@@ -0,0 +1,89 @@
+import torch
+import torch.nn.functional as F
+from tokenizers import Tokenizer
+
+# =========================
+# Juicy variables
+# =========================
+CHECKPOINT_PATH = "ckpt.pt"
+TOKENIZER_PATH  = "tokenizer/tokenizer.json"
+SEQ_LEN         = 128
+DEVICE          = "cuda" if torch.cuda.is_available() else "cpu"
+
+# =========================
+# Load tokenizer
+# =========================
+tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+pad_id = tokenizer.token_to_id("[PAD]")
+
+def encode_sentences(sentences):
+    ids = []
+    for s in sentences:
+        tokens = tokenizer.encode(s).ids
+        tokens = tokens[:SEQ_LEN]
+        tokens += [pad_id] * (SEQ_LEN - len(tokens))
+        ids.append(tokens)
+    return torch.tensor(ids, dtype=torch.long, device=DEVICE)
+
+# =========================
+# Model (must match training definition)
+# =========================
+class TransformerEncoder(torch.nn.Module):
+    def __init__(self, vocab_size, model_dim=256, num_layers=6, num_heads=4, ff_dim=1024, seq_len=128):
+        super().__init__()
+        self.token_emb = torch.nn.Embedding(vocab_size, model_dim)
+        self.pos_emb   = torch.nn.Embedding(seq_len, model_dim)
+
+        encoder_layer = torch.nn.TransformerEncoderLayer(
+            d_model=model_dim,
+            nhead=num_heads,
+            dim_feedforward=ff_dim,
+            activation="gelu",
+            batch_first=True
+        )
+        self.encoder = torch.nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.norm = torch.nn.LayerNorm(model_dim)
+
+    def forward(self, x):
+        positions = torch.arange(0, x.size(1), device=x.device).unsqueeze(0)
+        h = self.token_emb(x) + self.pos_emb(positions)
+        h = self.encoder(h)
+        h = self.norm(h)
+        return h.mean(dim=1)  # pooled embedding
+
+# =========================
+# Load checkpoint
+# =========================
+VOCAB_SIZE = tokenizer.get_vocab_size()
+model = TransformerEncoder(vocab_size=VOCAB_SIZE).to(DEVICE)
+model.load_state_dict(torch.load(CHECKPOINT_PATH, map_location=DEVICE))
+model.eval()
+
+print("[INFO] Model loaded.")
+
+# =========================
+# Test sentences
+# =========================
+sentences = [
+    "The quick brown fox jumps over the lazy dog.",
+    "Neural networks are changing artificial intelligence.",
+    "I love eating pizza on weekends.",
+    "Quantum physics is hard but fascinating.",
+]
+
+inputs = encode_sentences(sentences)
+with torch.no_grad():
+    embeddings = model(inputs)
+
+# Normalize for cosine sim
+embeddings = F.normalize(embeddings, dim=1)
+
+print("\nEmbeddings:")
+for s, e in zip(sentences, embeddings):
+    print(f"{s}\n -> {e[:5].cpu().numpy()}...")  # show first 5 dims
+
+print("\nCosine similarities:")
+sims = embeddings @ embeddings.T
+for i in range(len(sentences)):
+    row = ["{:.2f}".format(x.item()) for x in sims[i]]
+    print(f"{i}: {row}")