Spaces:
Sleeping
Sleeping
File size: 1,747 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 | import torch
from torch import nn
class AberLanguageModel(nn.Module):
def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, dropout):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embed_dim)
self.gru = nn.GRU(
input_size=embed_dim,
hidden_size=hidden_dim,
num_layers=num_layers,
dropout=dropout if num_layers > 1 else 0.0,
batch_first=True,
)
self.dropout = nn.Dropout(dropout)
self.head = nn.Linear(hidden_dim, vocab_size)
def forward(self, idx, hidden=None, targets=None):
emb = self.embedding(idx)
out, hidden = self.gru(emb, hidden)
out = self.dropout(out)
logits = self.head(out)
loss = None
if targets is not None:
loss = nn.functional.cross_entropy(
logits.reshape(-1, logits.size(-1)),
targets.reshape(-1),
)
return logits, hidden, loss
def generate(self, idx, max_new_tokens, eos_id, temperature=1.0, top_k=8):
hidden = None
for _ in range(max_new_tokens):
logits, hidden, _ = self(idx[:, -1:], hidden)
next_logits = logits[:, -1, :] / max(temperature, 1e-4)
if top_k is not None and top_k > 0:
values, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
next_logits[next_logits < values[:, [-1]]] = float("-inf")
probs = torch.softmax(next_logits, dim=-1)
next_token = torch.multinomial(probs, num_samples=1)
idx = torch.cat([idx, next_token], dim=1)
if int(next_token.item()) == eos_id:
break
return idx
|