Text Generation
Transformers
ONNX
Safetensors
GGUF
Transformers.js
English
gpt2
humor
parody
tiny
chat
lm-studio
ollama
llama.cpp
gpt
meme
funny
joke
satire
ram
ram-shortage
tiny-model
small-model
micro-model
toy-model
trained-from-scratch
custom-tokenizer
synthetic-data
conversational
chatbot
cpu
low-resource
text-generation-inference
| import math | |
| import time | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import GPT2LMHeadModel, GPT2TokenizerFast | |
| from config import BATCH_SIZE, LEARNING_RATE, STEPS | |
| REPLY_PREFIX = "RAM-ALTMAN:" | |
| HOOK_TOKENS = 8 | |
| HOOK_WEIGHT = 4.0 | |
| LOG_EVERY = 250 | |
| def read_segments(data_file: str) -> tuple[list[str], list[bool]]: | |
| texts, is_reply = [], [] | |
| with open(data_file, encoding="utf-8") as f: | |
| for line in f: | |
| if line.startswith(REPLY_PREFIX): | |
| texts += [REPLY_PREFIX, line.removeprefix(REPLY_PREFIX)] | |
| is_reply += [False, True] | |
| else: | |
| texts.append(line) | |
| is_reply.append(False) | |
| return texts, is_reply | |
| def reply_weights(length: int) -> list[float]: | |
| # The hook weighs more than the rant | |
| if length <= HOOK_TOKENS: | |
| return [1.0] * length | |
| return [HOOK_WEIGHT] * HOOK_TOKENS + [1.0] * (length - HOOK_TOKENS) | |
| def load_token_ids(hf_tokenizer: GPT2TokenizerFast, data_file: str, device: str) -> tuple[torch.Tensor, torch.Tensor]: | |
| texts, is_reply = read_segments(data_file) | |
| ids, weights = [], [] | |
| for segment_ids, reply in zip(hf_tokenizer(texts)["input_ids"], is_reply): | |
| ids += segment_ids | |
| weights += reply_weights(len(segment_ids)) if reply else [0.0] * len(segment_ids) | |
| return torch.tensor(ids, device=device), torch.tensor(weights, device=device) | |
| def cosine_lr(step: int) -> float: | |
| return LEARNING_RATE * 0.5 * (1 + math.cos(math.pi * step / STEPS)) | |
| def random_batch(token_ids: torch.Tensor, weights: torch.Tensor, context_length: int) -> tuple[torch.Tensor, torch.Tensor]: | |
| starts = torch.randint(0, len(token_ids) - context_length - 1, (BATCH_SIZE,)) | |
| inputs = torch.stack([token_ids[i:i + context_length] for i in starts]) | |
| input_weights = torch.stack([weights[i:i + context_length] for i in starts]) | |
| return inputs, input_weights | |
| def weighted_loss(logits: torch.Tensor, inputs: torch.Tensor, input_weights: torch.Tensor) -> torch.Tensor: | |
| predictions = logits[:, :-1].reshape(-1, logits.size(-1)) | |
| targets = inputs[:, 1:].reshape(-1) | |
| target_weights = input_weights[:, 1:].reshape(-1) | |
| token_losses = F.cross_entropy(predictions, targets, reduction="none") | |
| return (token_losses * target_weights).sum() / target_weights.sum().clamp_min(1e-8) | |
| def train_loop(model: GPT2LMHeadModel, token_ids: torch.Tensor, weights: torch.Tensor, context_length: int) -> None: | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01) | |
| start = time.time() | |
| for step in range(1, STEPS + 1): | |
| for group in optimizer.param_groups: | |
| group["lr"] = cosine_lr(step) | |
| inputs, input_weights = random_batch(token_ids, weights, context_length) | |
| loss = weighted_loss(model(input_ids=inputs).logits, inputs, input_weights) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) | |
| optimizer.step() | |
| if step == 1 or step % LOG_EVERY == 0: | |
| print(f"step {step}/{STEPS} - loss {loss.item():.3f} <> {time.time() - start:.0f}s", flush=True) | |