Spaces:
Sleeping
Sleeping
File size: 1,517 Bytes
f907cd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import random
import torch
from .data import build_training_text
from .model import AberLanguageModel
from .tokenizer import WordTokenizer
def set_seed(seed: int):
random.seed(seed)
torch.manual_seed(seed)
def create_model_and_tokenizer(config, extra_text=""):
text = build_training_text(extra_text)
tokenizer = WordTokenizer().fit(text)
encoded = tokenizer.encode(text, add_bos=True, add_eos=True)
encoded = torch.tensor(encoded, dtype=torch.long)
model = AberLanguageModel(
vocab_size=tokenizer.vocab_size,
embed_dim=config.embed_dim,
hidden_dim=config.hidden_dim,
num_layers=config.num_layers,
dropout=config.dropout,
)
return model, tokenizer, encoded
def build_batch(encoded, seq_len, batch_size):
max_start = max(1, len(encoded) - seq_len - 1)
starts = torch.randint(0, max_start, (batch_size,))
x = torch.stack([encoded[start : start + seq_len] for start in starts])
y = torch.stack([encoded[start + 1 : start + seq_len + 1] for start in starts])
return x, y
def train_model(model, encoded, config, steps):
optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
model.train()
losses = []
for _ in range(steps):
xb, yb = build_batch(encoded, config.seq_len, config.batch_size)
_, _, loss = model(xb, targets=yb)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(float(loss.item()))
return losses
|