File size: 1,554 Bytes
a148a7d | 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 53 54 55 56 57 58 59 60 61 62 63 64 | import torch
from eztokenizer import EzTokenizer
import torch.nn as nn
tokenizer = EzTokenizer()
dataset = """
blub blub blubs blubby
ooh food
bloob blub blip blop
WAIT IS THAT A SHARK NO PLEZ HELP ME
Mama? i want food
beep blab bloop blob boloob food
no shark yipe!!!!!!!!!!!
PLEZ MR SHARK NO EAT ME TANKZ
ooh i like seaweed em nom nom
ooh coral reef wow
yummi coral reef nom nom chomp blub
"""
tokenizer.train(dataset)
enc = tokenizer.encode(dataset)
def shift(tokens, context_len: int):
x_sequences = []
y_sequences = []
for position in range(len(tokens) - context_len):
context = tokens[position:position + context_len]
x_sequences.append(context)
y_context = tokens[position+1: position+context_len+1]
y_sequences.append(y_context)
return x_sequences, y_sequences
x, y = shift(enc, 32)
x = torch.tensor(x)
y = torch.tensor(y)
embed_dim = 6
num_heads = 3
embedding = torch.nn.Embedding(len(tokenizer.vocab) + 1, 6)
attention = torch.nn.MultiheadAttention(embed_dim, num_heads, dropout=0.0, bias=True)
output = torch.nn.Linear(6, len(tokenizer.vocab) + 1, bias=True)
loss = torch.nn.CrossEntropyLoss()
class Fih(nn.Module):
def __init__(self):
super().__init__()
self.embedding = embedding
self.attention = attention
self.output = output
def forward_pass(self, x):
input_tokens = x
em_x = self.embedding(x)
fihs_attention = self.attention(em_x, em_x, em_x)[0]
fih_output = self.output(fihs_attention)
return fih_output |