import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader, TensorDataset import numpy as np import os import re from collections import Counter # ------------------------------ # 1. Prepare data from books # ------------------------------ def load_books(books_dir): """Load and tokenize all markdown books.""" text = "" for filename in sorted(os.listdir(books_dir)): if filename.endswith(".md"): with open(os.path.join(books_dir, filename), "r", encoding="utf-8") as f: text += f.read() + "\n" return text def tokenize(text): """Simple tokenization: lowercase, keep only alphabetic words.""" # Remove markdown headers and special chars text = re.sub(r'#+', '', text) text = re.sub(r'[*_~`]', '', text) # Extract words (only alphabetic) words = re.findall(r'[a-z]+', text.lower()) return words books_dir = "/home/per/Documents/lag word ai/books" raw_text = load_books(books_dir) words = tokenize(raw_text) print(f"Total words: {len(words)}") # Build vocabulary (filter rare words) min_freq = 5 word_counts = Counter(words) vocab = [w for w, c in word_counts.items() if c >= min_freq] vocab = sorted(vocab) # deterministic order word2idx = {w: i for i, w in enumerate(vocab)} idx2word = {i: w for i, w in enumerate(vocab)} vocab_size = len(vocab) print(f"Vocabulary size (min_freq={min_freq}): {vocab_size}") # Convert words to indices (use for rare words) unk_idx = vocab_size word_indices = [word2idx.get(w, unk_idx) for w in words] vocab_size += 1 # include print(f"Vocabulary with : {vocab_size}") # Compute word heights based on frequency (normalized) max_count = max(word_counts[w] for w in vocab) if vocab else 1 word_heights = {} for w in vocab: # Normalize frequency to [0.1, 0.9] range freq = word_counts[w] / max_count word_heights[word2idx[w]] = 0.1 + 0.8 * freq word_heights[unk_idx] = 0.05 # low height for unknown words # Build sliding windows (lag = n) n = 3 # use previous 3 words to predict the next # Subsample for faster training (use every k-th sequence) subsample = 10 # change to 1 for full dataset def create_sequences(word_indices, word_heights, n, subsample=1): X_words, X_heights, y_words = [], [], [] for i in range(0, len(word_indices) - n, subsample): in_words = word_indices[i:i+n] in_heights = [word_heights[w] for w in in_words] target = word_indices[i+n] X_words.append(in_words) X_heights.append(in_heights) y_words.append(target) return (torch.tensor(X_words, dtype=torch.long), torch.tensor(X_heights, dtype=torch.float), torch.tensor(y_words, dtype=torch.long)) X_words, X_heights, y_words = create_sequences(word_indices, word_heights, n, subsample) print(f"X_words shape: {X_words.shape}") print(f"X_heights shape: {X_heights.shape}") print(f"y_words shape: {y_words.shape}") dataset = TensorDataset(X_words, X_heights, y_words) dataloader = DataLoader(dataset, batch_size=10, shuffle=False) # ------------------------------ # 2. Define the model # ------------------------------ class WordPredictorWithHeight(nn.Module): def __init__(self, vocab_size, n, embed_dim=32, hidden_dim=64): super().__init__() self.n = n self.word_embed = nn.Embedding(vocab_size, embed_dim) # We will concatenate word embedding (embed_dim) + height scalar (1) # So input size to LSTM = embed_dim + 1 self.lstm = nn.LSTM(input_size=embed_dim + 1, hidden_size=hidden_dim, batch_first=True) self.fc = nn.Linear(hidden_dim, vocab_size) def forward(self, word_seq, height_seq): # word_seq: (batch, n) # height_seq: (batch, n) batch_size = word_seq.size(0) # Embed words word_emb = self.word_embed(word_seq) # (batch, n, embed_dim) # Add height as an extra feature channel height_seq = height_seq.unsqueeze(-1) # (batch, n, 1) combined = torch.cat([word_emb, height_seq], dim=-1) # (batch, n, embed_dim+1) # LSTM lstm_out, _ = self.lstm(combined) # (batch, n, hidden_dim) # Use only the last time step's output last_out = lstm_out[:, -1, :] # (batch, hidden_dim) logits = self.fc(last_out) # (batch, vocab_size) return logits # ------------------------------ # 3. Training loop # ------------------------------ model = WordPredictorWithHeight(vocab_size, n) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=0.001) epochs = 200 for epoch in range(epochs): total_loss = 0.0 for batch_words, batch_heights, batch_targets in dataloader: optimizer.zero_grad() logits = model(batch_words, batch_heights) loss = criterion(logits, batch_targets) loss.backward() optimizer.step() total_loss += loss.item() avg_loss = total_loss / len(dataloader) print(f"Epoch {epoch+1:3d}, Loss: {avg_loss:.4f}") # ------------------------------ # 4. Evaluation: predict next word # ------------------------------ def predict_next(model, word_seq, height_seq, idx2word, unk_idx, topk=5): """word_seq: list of n word indices, height_seq: list of n heights""" model.eval() with torch.no_grad(): words_t = torch.tensor([word_seq], dtype=torch.long) heights_t = torch.tensor([height_seq], dtype=torch.float) logits = model(words_t, heights_t) probs = torch.softmax(logits, dim=1) topk_probs, topk_indices = torch.topk(probs, topk, dim=1) results = [] for i in range(topk): idx = topk_indices[0, i].item() prob = topk_probs[0, i].item() word = idx2word.get(idx, "") results.append((word, prob)) return results # Example: use last 3 words from the text test_words_str = words[-n:] print(f"\nInput words: {test_words_str}") test_word_indices = [word2idx.get(w, unk_idx) for w in test_words_str] test_heights = [word_heights[w] for w in test_word_indices] predictions = predict_next(model, test_word_indices, test_heights, idx2word, unk_idx) print(f"Top predictions for next word:") for word, prob in predictions: print(f" {word}: {prob:.3f}")