import os
import json
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
# =========================
# SETTINGS
# =========================
KNOWLEDGE_FILE = "knowledge.txt"
MODEL_FILE = "model.pt"
VOCAB_FILE = "vocab.json"
OPTIMIZER_FILE = "optimizer.pt"
BATCH_SIZE = 128
CONTEXT_SIZE = 128
EPOCHS_PER_RUN = 50
LEARNING_RATE = 0.000001
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ~1 million parameter model
EMBED_SIZE = 128
NUM_HEADS = 4
NUM_LAYERS = 4
FF_SIZE = 512
# =========================
# TOKENIZER
# =========================
def build_vocab(text):
chars = sorted(list(set(text)))
stoi = {ch: i for i, ch in enumerate(chars)}
itos = {i: ch for i, ch in enumerate(chars)}
return stoi, itos
def save_vocab(stoi):
with open(VOCAB_FILE, "w", encoding="utf-8") as f:
json.dump(stoi, f)
def load_vocab():
with open(VOCAB_FILE, "r", encoding="utf-8") as f:
stoi = json.load(f)
itos = {int(v): k for k, v in zip(stoi.keys(), stoi.values())}
return stoi, itos
# =========================
# DATASET
# =========================
class CharDataset(Dataset):
def __init__(self, data, block_size):
self.data = data
self.block_size = block_size
def __len__(self):
return len(self.data) - self.block_size - 1
def __getitem__(self, idx):
x = self.data[idx:idx+self.block_size]
y = self.data[idx+1:idx+self.block_size+1]
return torch.tensor(x), torch.tensor(y)
# =========================
# MODEL
# =========================
class GPTBlock(nn.Module):
def __init__(self, embed_size, heads, ff_size):
super().__init__()
self.attn = nn.MultiheadAttention(
embed_size,
heads,
batch_first=True
)
self.ln1 = nn.LayerNorm(embed_size)
self.ln2 = nn.LayerNorm(embed_size)
self.ff = nn.Sequential(
nn.Linear(embed_size, ff_size),
nn.ReLU(),
nn.Linear(ff_size, embed_size)
)
def forward(self, x):
seq_len = x.size(1)
mask = torch.triu(
torch.ones(seq_len, seq_len, device=x.device),
diagonal=1
).bool()
attn_out, _ = self.attn(
x, x, x,
attn_mask=mask
)
x = self.ln1(x + attn_out)
x = self.ln2(x + self.ff(x))
return x
class GPTModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.token_emb = nn.Embedding(vocab_size, EMBED_SIZE)
self.pos_emb = nn.Embedding(CONTEXT_SIZE, EMBED_SIZE)
self.blocks = nn.ModuleList([
GPTBlock(
EMBED_SIZE,
NUM_HEADS,
FF_SIZE
)
for _ in range(NUM_LAYERS)
])
self.ln = nn.LayerNorm(EMBED_SIZE)
self.head = nn.Linear(EMBED_SIZE, vocab_size)
def forward(self, idx):
B, T = idx.shape
tok = self.token_emb(idx)
pos = self.pos_emb(
torch.arange(T, device=idx.device)
)
x = tok + pos
for block in self.blocks:
x = block(x)
x = self.ln(x)
return self.head(x)
# =========================
# PARAMETER COUNT
# =========================
def count_parameters(model):
return sum(
p.numel()
for p in model.parameters()
if p.requires_grad
)
# =========================
# TRAIN
# =========================
def train_epoch(model, loader, optimizer):
model.train()
total_batches = len(loader)
for batch_idx, (x, y) in enumerate(loader):
x = x.to(DEVICE)
y = y.to(DEVICE)
optimizer.zero_grad()
logits = model(x)
loss = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
y.reshape(-1)
)
loss.backward()
optimizer.step()
percent = (batch_idx + 1) / total_batches * 100
print(
f"\rBatch {batch_idx+1}/{total_batches} "
f"({percent:.1f}%) "
f"Loss: {loss.item():.4f}",
end=""
)
print()
# =========================
# GENERATE
# =========================
@torch.no_grad()
def generate(model, prompt, stoi, itos,
max_new_chars=512,
temperature=0.8):
model.eval()
ids = [stoi.get(ch, 0) for ch in prompt]
x = torch.tensor(
[ids],
dtype=torch.long,
device=DEVICE
)
for _ in range(max_new_chars):
x_cond = x[:, -CONTEXT_SIZE:]
logits = model(x_cond)
logits = logits[:, -1, :] / temperature
probs = F.softmax(logits, dim=-1)
next_id = torch.multinomial(
probs,
num_samples=1
)
x = torch.cat(
[x, next_id],
dim=1
)
result = "".join(
itos[int(i)]
for i in x[0].tolist()
)
return result
# =========================
# MAIN
# =========================
def main():
if not os.path.exists(KNOWLEDGE_FILE):
print("knowledge.txt not found.")
return
with open(
KNOWLEDGE_FILE,
"r",
encoding="utf-8"
) as f:
text = f.read()
if len(text) < CONTEXT_SIZE + 2:
print("knowledge.txt too small.")
return
if os.path.exists(VOCAB_FILE):
with open(
VOCAB_FILE,
"r",
encoding="utf-8"
) as f:
stoi = json.load(f)
itos = {
int(v): k
for k, v in zip(
stoi.keys(),
stoi.values()
)
}
else:
stoi, itos = build_vocab(text)
save_vocab(stoi)
vocab_size = len(stoi)
encoded = [
stoi[ch]
for ch in text
if ch in stoi
]
dataset = CharDataset(
encoded,
CONTEXT_SIZE
)
loader = DataLoader(
dataset,
batch_size=BATCH_SIZE,
shuffle=True,
drop_last=True
)
model = GPTModel(vocab_size)
if os.path.exists(MODEL_FILE):
print("Loading existing model...")
model.load_state_dict(
torch.load(
MODEL_FILE,
map_location="cpu"
)
)
model = model.to(DEVICE)
params = count_parameters(model)
print("\n========================")
print(f"Device: {DEVICE}")
print(f"Vocabulary Size: {vocab_size}")
print(f"Model Parameters: {params:,}")
print("========================\n")
optimizer = torch.optim.AdamW(
model.parameters(),
lr=LEARNING_RATE
)
# Load optimizer state if it exists
if os.path.exists(OPTIMIZER_FILE):
print("Loading existing optimizer state...")
optimizer.load_state_dict(
torch.load(
OPTIMIZER_FILE,
map_location=DEVICE
)
)
for epoch in range(EPOCHS_PER_RUN):
print(
f"Epoch {epoch+1}/{EPOCHS_PER_RUN}"
)
train_epoch(
model,
loader,
optimizer
)
torch.save(
model.state_dict(),
MODEL_FILE
)
torch.save(
optimizer.state_dict(),
OPTIMIZER_FILE
)
# Create INT8 quantized version
quantized_model = torch.quantization.quantize_dynamic(
model.cpu(),
{nn.Linear},
dtype=torch.qint8
)
torch.save(
quantized_model.state_dict(),
"model_int8.pt"
)
# Move original model back to the selected device
model = model.to(DEVICE)
print("\nModel saved.")
print("Optimizer state saved.")
print("INT8 model saved as model_int8.pt")
while True:
prompt = input("\nYou: ")
if prompt.lower() in [
"quit",
"exit",
"q"
]:
break
reply = generate(
model,
prompt,
stoi,
itos,
max_new_chars=512
)
print("\nAI:")
print(reply)
if __name__ == "__main__":
main()