import os
import math
import json
import torch
import torch.nn as nn
import torch.nn.functional as F
# ============================================================
# CHARACTER-LEVEL TOKENIZER
# ============================================================
class EncodedText:
def __init__(self, ids):
self.ids = ids
class CharacterTokenizer:
SPECIAL_TOKENS = ["[UNK]", "[PAD]", "[BOS]", "[EOS]"]
def __init__(self, vocab):
self.vocab = vocab
self.id_to_token = {
token_id: token
for token, token_id in vocab.items()
}
self.unk_id = self.vocab["[UNK]"]
@classmethod
def from_text(cls, text):
vocab = {
token: index
for index, token in enumerate(cls.SPECIAL_TOKENS)
}
for character in sorted(set(text)):
if character not in vocab:
vocab[character] = len(vocab)
return cls(vocab)
@classmethod
def from_file(cls, filename):
with open(filename, "r", encoding="utf-8") as file:
saved_data = json.load(file)
if saved_data.get("type") != "character":
raise ValueError("Existing tokenizer is not character-level.")
return cls(saved_data["vocab"])
def save(self, filename):
with open(filename, "w", encoding="utf-8") as file:
json.dump(
{
"type": "character",
"vocab": self.vocab
},
file,
ensure_ascii=False,
indent=2
)
def encode(self, text, add_special_tokens=False):
return EncodedText([
self.vocab.get(character, self.unk_id)
for character in text
])
def decode(self, ids, skip_special_tokens=False):
special_ids = {
self.vocab[token]
for token in self.SPECIAL_TOKENS
}
characters = []
for token_id in ids:
token_id = int(token_id)
if skip_special_tokens and token_id in special_ids:
continue
characters.append(
self.id_to_token.get(token_id, "[UNK]")
)
return "".join(characters)
def token_to_id(self, token):
return self.vocab.get(token)
def get_vocab_size(self):
return len(self.vocab)
# ============================================================
# SETTINGS
# ============================================================
KNOWLEDGE_FILE = "knowledge.txt"
MODEL_FILE = "model.pt"
TOKENIZER_FILE = "tokenizer.json"
VOCAB_SIZE = 8000
EMBED_SIZE = 192
NUM_HEADS = 4
NUM_LAYERS = 4
FF_SIZE = 768
CONTEXT_SIZE = 256
BATCH_SIZE = 16
TRAINING_STEPS = 1000
LEARNING_RATE = 0.0003
DROPOUT = 0.1
SAVE_INTERVAL = 100
MAX_NEW_TOKENS = 2000
TEMPERATURE = 0.8
TOP_K = 40
DEVICE = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
# ============================================================
# TOKENIZER
# ============================================================
def create_tokenizer():
if not os.path.exists(KNOWLEDGE_FILE):
raise FileNotFoundError(KNOWLEDGE_FILE)
with open(
KNOWLEDGE_FILE,
"r",
encoding="utf-8"
) as file:
text = file.read()
if not text:
raise ValueError("knowledge.txt is empty.")
if os.path.exists(TOKENIZER_FILE):
print("Loading existing character-level tokenizer...")
try:
tokenizer = CharacterTokenizer.from_file(TOKENIZER_FILE)
missing_characters = set(text) - set(tokenizer.vocab)
if not missing_characters:
return tokenizer
print(
"Training text contains new characters. "
"Recreating tokenizer..."
)
except (ValueError, KeyError, json.JSONDecodeError):
print(
"Existing tokenizer is incompatible. "
"Creating a character-level tokenizer..."
)
print("Creating character-level tokenizer...")
tokenizer = CharacterTokenizer.from_text(text)
tokenizer.save(TOKENIZER_FILE)
print("Character-level tokenizer created.")
return tokenizer
# ============================================================
# MODEL
# ============================================================
class CausalAttention(nn.Module):
def __init__(self):
super().__init__()
self.heads = NUM_HEADS
self.head_size = EMBED_SIZE // NUM_HEADS
self.qkv = nn.Linear(
EMBED_SIZE,
EMBED_SIZE * 3
)
self.projection = nn.Linear(
EMBED_SIZE,
EMBED_SIZE
)
self.dropout = nn.Dropout(DROPOUT)
def forward(self, x):
batch, sequence, channels = x.shape
q, k, v = self.qkv(x).chunk(3, dim=-1)
q = q.view(
batch,
sequence,
self.heads,
self.head_size
).transpose(1, 2)
k = k.view(
batch,
sequence,
self.heads,
self.head_size
).transpose(1, 2)
v = v.view(
batch,
sequence,
self.heads,
self.head_size
).transpose(1, 2)
attention = F.scaled_dot_product_attention(
q,
k,
v,
is_causal=True,
dropout_p=self.dropout.p if self.training else 0.0
)
attention = attention.transpose(
1, 2
).contiguous().view(
batch,
sequence,
channels
)
return self.dropout(
self.projection(attention)
)
class FeedForward(nn.Module):
def __init__(self):
super().__init__()
self.network = nn.Sequential(
nn.Linear(EMBED_SIZE, FF_SIZE),
nn.GELU(),
nn.Linear(FF_SIZE, EMBED_SIZE),
nn.Dropout(DROPOUT)
)
def forward(self, x):
return self.network(x)
class TransformerBlock(nn.Module):
def __init__(self):
super().__init__()
self.norm1 = nn.LayerNorm(EMBED_SIZE)
self.attention = CausalAttention()
self.norm2 = nn.LayerNorm(EMBED_SIZE)
self.feed_forward = FeedForward()
def forward(self, x):
x = x + self.attention(
self.norm1(x)
)
x = x + self.feed_forward(
self.norm2(x)
)
return x
class GPTModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.token_embedding = nn.Embedding(
vocab_size,
EMBED_SIZE
)
self.position_embedding = nn.Embedding(
CONTEXT_SIZE,
EMBED_SIZE
)
self.dropout = nn.Dropout(DROPOUT)
self.blocks = nn.Sequential(
*[
TransformerBlock()
for _ in range(NUM_LAYERS)
]
)
self.final_norm = nn.LayerNorm(EMBED_SIZE)
self.output = nn.Linear(
EMBED_SIZE,
vocab_size,
bias=False
)
# Share input and output embeddings.
self.output.weight = self.token_embedding.weight
self.apply(self.initialize_weights)
def initialize_weights(self, module):
if isinstance(module, nn.Linear):
nn.init.normal_(
module.weight,
mean=0.0,
std=0.02
)
if module.bias is not None:
nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
nn.init.normal_(
module.weight,
mean=0.0,
std=0.02
)
def forward(self, x, targets=None):
batch, sequence = x.shape
positions = torch.arange(
sequence,
device=x.device
)
token_embeddings = self.token_embedding(x)
position_embeddings = self.position_embedding(
positions
)
x = token_embeddings + position_embeddings
x = self.dropout(x)
x = self.blocks(x)
x = self.final_norm(x)
logits = self.output(x)
loss = None
if targets is not None:
loss = F.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1)
)
return logits, loss
@torch.no_grad()
def generate(
self,
input_ids,
max_new_tokens=MAX_NEW_TOKENS,
temperature=TEMPERATURE,
top_k=TOP_K
):
self.eval()
for _ in range(max_new_tokens):
input_context = input_ids[
:, -CONTEXT_SIZE:
]
logits, _ = self(input_context)
logits = logits[:, -1, :]
logits = logits / max(temperature, 0.01)
if top_k > 0:
k = min(top_k, logits.size(-1))
values, _ = torch.topk(
logits,
k
)
threshold = values[:, [-1]]
logits = logits.masked_fill(
logits < threshold,
float("-inf")
)
probabilities = F.softmax(
logits,
dim=-1
)
next_token = torch.multinomial(
probabilities,
num_samples=1
)
input_ids = torch.cat(
(input_ids, next_token),
dim=1
)
return input_ids
# ============================================================
# SAVE AND LOAD
# ============================================================
def save_checkpoint(model, optimizer, step):
temporary_file = MODEL_FILE + ".tmp"
torch.save(
{
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"step": step,
"config": {
"vocab_size": model.token_embedding.num_embeddings,
"embed_size": EMBED_SIZE,
"num_heads": NUM_HEADS,
"num_layers": NUM_LAYERS,
"ff_size": FF_SIZE,
"context_size": CONTEXT_SIZE
}
},
temporary_file
)
os.replace(
temporary_file,
MODEL_FILE
)
print("Model saved.")
def load_checkpoint(model, optimizer):
if not os.path.exists(MODEL_FILE):
return 0
print("Loading existing model...")
checkpoint = torch.load(
MODEL_FILE,
map_location=DEVICE,
weights_only=False
)
expected_config = {
"vocab_size": model.token_embedding.num_embeddings,
"embed_size": EMBED_SIZE,
"num_heads": NUM_HEADS,
"num_layers": NUM_LAYERS,
"ff_size": FF_SIZE,
"context_size": CONTEXT_SIZE
}
if checkpoint["config"] != expected_config:
raise ValueError(
"Model configuration does not match the saved model. "
"Restore the original settings or use a new model filename."
)
model.load_state_dict(
checkpoint["model"]
)
if "optimizer" in checkpoint:
optimizer.load_state_dict(
checkpoint["optimizer"]
)
print("Existing model loaded.")
return checkpoint.get("step", 0)
# ============================================================
# LOAD TRAINING DATA
# ============================================================
def load_training_data(tokenizer):
if not os.path.exists(KNOWLEDGE_FILE):
raise FileNotFoundError(
"knowledge.txt was not found."
)
with open(
KNOWLEDGE_FILE,
"r",
encoding="utf-8"
) as file:
text = file.read()
if not text.strip():
raise ValueError(
"knowledge.txt is empty."
)
encoded = tokenizer.encode(text)
tokens = encoded.ids
eos_id = tokenizer.token_to_id("[EOS]")
tokens.append(eos_id)
if len(tokens) <= CONTEXT_SIZE:
raise ValueError(
"Not enough training tokens. "
"Add more text to knowledge.txt or reduce CONTEXT_SIZE."
)
return torch.tensor(
tokens,
dtype=torch.long
)
# ============================================================
# RANDOM TRAINING BATCHES
# ============================================================
def get_batch(data):
max_start = len(data) - CONTEXT_SIZE - 1
starts = torch.randint(
0,
max_start + 1,
(BATCH_SIZE,)
)
x = torch.stack([
data[i:i + CONTEXT_SIZE]
for i in starts.tolist()
])
y = torch.stack([
data[i + 1:i + CONTEXT_SIZE + 1]
for i in starts.tolist()
])
return (
x.to(DEVICE),
y.to(DEVICE)
)
# ============================================================
# TRAINING
# ============================================================
def train_model(model, optimizer, data, start_step):
model.train()
print("\nStarting training...")
for step in range(
start_step + 1,
start_step + TRAINING_STEPS + 1
):
x, y = get_batch(data)
optimizer.zero_grad(set_to_none=True)
_, loss = model(x, y)
loss.backward()
torch.nn.utils.clip_grad_norm_(
model.parameters(),
1.0
)
optimizer.step()
if step % 10 == 0:
print(
f"Step {step} | Loss: {loss.item():.4f}"
)
if step % SAVE_INTERVAL == 0:
save_checkpoint(
model,
optimizer,
step
)
save_checkpoint(
model,
optimizer,
start_step + TRAINING_STEPS
)
print("\nTraining completed.")
# ============================================================
# CHAT
# ============================================================
def chat(model, tokenizer):
print("\n================================")
print(" GPT CHAT")
print("================================")
print("Type /exit to quit.\n")
eos_id = tokenizer.token_to_id("[EOS]")
while True:
user_input = input("You: ").strip()
if user_input.lower() == "/exit":
break
if not user_input:
continue
prompt = (
"User: "
+ user_input
+ "\nAssistant:"
)
encoded = tokenizer.encode(
prompt,
add_special_tokens=False
)
input_ids = torch.tensor(
[encoded.ids],
dtype=torch.long,
device=DEVICE
)
output = model.generate(
input_ids,
max_new_tokens=MAX_NEW_TOKENS,
temperature=TEMPERATURE,
top_k=TOP_K
)
generated_ids = output[0][
len(encoded.ids):
].tolist()
if eos_id in generated_ids:
generated_ids = generated_ids[
:generated_ids.index(eos_id)
]
response = tokenizer.decode(
generated_ids,
skip_special_tokens=True
)
# Remove accidental continuation into another user turn.
if "\nUser:" in response:
response = response.split(
"\nUser:",
1
)[0]
print("\nAssistant:", response.strip(), "\n")
model.eval()
# ============================================================
# MAIN PROGRAM
# ============================================================
def main():
print("================================")
print(" SMALL GPT MODEL")
print("================================")
print("Device:", DEVICE)
tokenizer = create_tokenizer()
vocab_size = tokenizer.get_vocab_size()
print("Vocabulary size:", vocab_size)
model = GPTModel(
vocab_size
).to(DEVICE)
parameter_count = sum(
p.numel()
for p in model.parameters()
)
print(
"Parameters:",
f"{parameter_count:,}"
)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=LEARNING_RATE,
weight_decay=0.01
)
start_step = load_checkpoint(
model,
optimizer
)
data = load_training_data(
tokenizer
)
print(
"Training tokens:",
f"{len(data):,}"
)
print(
"Previous training steps:",
start_step
)
train_model(
model,
optimizer,
data,
start_step
)
chat(
model,
tokenizer
)
if __name__ == "__main__":
try:
main()
except KeyboardInterrupt:
print("\nProgram interrupted.")
except Exception as error:
print("\nERROR:", error)
input("\nPress Enter to close...")