[Python] Vocab size tensors index out of range

I am trying to create a tensor (vocab_size x vocab_size) for a llm, I am using the tiktoken module. But the problem is that when I create the tensor the indices are out of range. The vocab size I am using is the vocab in the text rather than the entire vocab size of tiktoken as the tensor would be much too large:

encoding = tk.get_encoding('p50k_base')
encoded_text = encoding.encode(text)
vocab_size = len(set(encoded_text))

data = torch.tensor(encoded_text, dtype=torch.long)

num = int(len(data) * 0.9)
train = data[:num]
test = data[num:]

block_size = 8
batch_size = 4

x = train[:block_size]
y = train[1:block_size+1]
for t in range(block_size):
context = x[:t+1]
target = y[t]

def get_batch(block_size, batch_size):
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([data[i:i+block_size] for i in ix])
y = torch.stack([data[i+1: i+block_size+1] for i in ix])
return x, y

xb, yb = get_batch(block_size, batch_size)

class BigramLanguageModel(nn.Module):
def __init__(self, vocab_size):
super().__init__()
self.token_embedding_table = nn.Embedding(vocab_size, vocab_size)

def forward(self, idx, targets=None):
print(idx)
logits = self.token_embedding_table(idx)

if targets is None:
loss = None
else:
B,T,C = logits.shape
logits = logits.view(B*T, C)
targets = targets.view(B*T)
loss = F.cross_entropy(logits, targets)

return logits, loss
def generate(self, idx, max_new_tokens):
logits, loss = self(idx)
logits = logits[:, -1, :]
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
idx = torch.cat((idx, idx_next), dim=1,)

return idx


I tried to use the function n_vocab but this returns the whole vocab size for tiktoken rather than my text and creates a massive tensor. But it does solve the index problem.

Continue reading...