add: v0

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+output/

app.py

@@ -0,0 +1,45 @@
+import torch
+import gradio as gr
+from model import CharacterLevelTokenizer, PotterGPT, Config
+
+class GradioApp():
+    def __init__(self):
+        # Set up configuration and data
+        self.model_path = 'potterGPT/potterGPT.pth'
+        with open('data/harry_potter_data', 'r', encoding='utf-8') as f:
+            data = f.read()
+
+        self.tokenizer = CharacterLevelTokenizer(data)
+        self.lm = PotterGPT(Config)
+        state_dict = torch.load(self.model_path, map_location='cpu')
+        self.lm.load_state_dict(state_dict)
+
+    def launch(self):
+        # Define Gradio interface without a clear button
+        with gr.Blocks() as demo:
+            gr.Markdown("# potterGPT v0")
+            gr.Markdown("Click the button to generate a text prompt using the potterGPT model.")
+            
+            generate_button = gr.Button("Generate")
+            output_text = gr.Textbox(label="Generated Text")
+            
+            generate_button.click(self.generate_text, inputs=None, outputs=output_text)
+        
+        demo.launch()
+
+    def generate_text(self, input=None):
+        """Generate text using the trained model."""
+        generated_texts = []
+        for length in [1000]:
+            generated = self.lm.generate(
+                torch.zeros((1,1),dtype=torch.long,device='cpu') + 61, # initial context 0, 61 is \n
+                total=length
+            )
+            generated = self.tokenizer.decode(generated[0].cpu().numpy())
+            text = f'generated ({length} tokens)\n{"="*50}\n{generated}\n{"="*50}\n\n'
+            generated_texts.append(text)
+        return generated_texts[0]
+
+if __name__ == '__main__':
+    app = GradioApp()
+    app.launch()

build_tokenizer.py

@@ -0,0 +1,22 @@
+from pathlib import Path
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+import tokenizers.pre_tokenizers as pre_tokenizers
+import tokenizers.processors as processors
+import tokenizers.decoders as decoders
+from tokenizers.trainers import BpeTrainer
+
+if __name__ == '__main__':
+
+    tokenizer_path = Path('tokenizer/')
+    tokenizer_path.mkdir(exist_ok=True)
+
+    tokenizer = Tokenizer(BPE())
+
+    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
+    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)
+    tokenizer.decoder = decoders.ByteLevel()
+    trainer = BpeTrainer(special_tokens=['<|endoftext|>'], min_frequency=2)
+
+    tokenizer.train(['data/harry_potter_data'],trainer)
+    tokenizer.save(str(tokenizer_path / 'potter.json'))

inference.py

@@ -0,0 +1,35 @@
+import torch
+import os
+from model import PotterGPT, Config, CharacterLevelTokenizer
+from tokenizers import Tokenizer
+from dataclasses import dataclass
+
+model_path = 'potterGPT/potterGPT.pth'
+with open('data/harry_potter_data', 'r', encoding='utf-8') as f:
+    data = f.read()
+
+
+tokenizer = CharacterLevelTokenizer(data)
+
+
+
+lm = PotterGPT(Config)
+state_dict = torch.load(model_path, map_location='cpu')
+lm.load_state_dict(state_dict)
+
+generated_texts = []
+for length in [1000]:
+    generated = lm.generate(
+    torch.zeros((1,1),dtype=torch.long,device='cpu') + 61, # initial context 0, 61 is \n
+    total=length
+)
+    generated = tokenizer.decode(generated[0].cpu().numpy())
+    text=f'generated ({length} tokens)\n{"="*50}\n{generated}\n{"="*50}\n\n'
+    generated_texts.append(text)
+
+print(generated_texts[0])
+
+os.makedirs('output', exist_ok=True)
+with open('output/generated.txt', 'w+') as f:
+    for text in generated_texts:
+        f.write(text)

model/__init__.py

@@ -0,0 +1,7 @@
+from .config import Config
+from .transformer_block import TransformerBlock
+from .feed_forward import FeedForward
+from .multi_head_attention import MultiHeadAttention
+from .single_attention_head import AttentionHead
+from .model import PotterGPT
+from .tokenizer import CharacterLevelTokenizer

model/config.py

@@ -0,0 +1,27 @@
+import torch
+from .tokenizer import CharacterLevelTokenizer
+from dataclasses import dataclass
+
+with open('data/harry_potter_data', 'r', encoding='utf-8') as f:
+    data = f.read()
+
+@dataclass
+class Config:
+    tokenizer = CharacterLevelTokenizer(data)
+    block_size = 256 # context-length
+    batch_size = 64 # mini-batch size
+    vocab_size = tokenizer.VOCAB_SIZE
+    n_embed = 256
+    n_heads = 8
+    head_size =n_embed //n_heads # computes to 384/6=64 or 128/4=32 or 256/8
+        
+    n_layers = 3
+        
+    train_iters = 10_000
+    val_iters = 1000
+    lr = 3e-4
+        
+    attn_dropout = 0.1
+    block_dropout = 0.1
+        
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'

model/feed_forward.py

@@ -0,0 +1,14 @@
+import torch
+import torch.nn as nn
+
+class FeedForward(nn.Module):
+    def __init__(self, Config):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(Config.n_embed,Config.n_embed * 4),
+            nn.ReLU(),
+            nn.Linear(Config.n_embed * 4, Config.n_embed), # projection
+            nn.Dropout(Config.block_dropout)
+        )
+    def forward(self,x):
+        return self.net(x)

model/model.py

@@ -0,0 +1,47 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .transformer_block import TransformerBlock
+from .config import Config
+
+class PotterGPT(nn.Module):
+    def __init__(self,Config):
+        super().__init__()
+        
+        self.n_embed = Config.n_embed
+        self.block_size = Config.block_size
+        
+        self.token_embedding_table = nn.Embedding(Config.vocab_size,self.n_embed)
+        self.pos_embedding_table = nn.Embedding(self.block_size, self.n_embed)
+        
+        self.blocks = nn.Sequential(
+            *[TransformerBlock(Config)]*Config.n_layers,
+            nn.LayerNorm(self.n_embed)
+        )
+
+        self.lm_head = nn.Linear(self.n_embed,Config.vocab_size)
+        
+    def forward(self,idx):
+        
+        B,T = idx.shape
+        
+        token_embs = self.token_embedding_table(idx)
+        pos_embs = self.pos_embedding_table(torch.arange(T,device=Config.device))
+        
+        
+        x = token_embs + pos_embs
+        x = self.blocks(x)
+        logits = self.lm_head(x)
+        
+        return logits
+
+        
+    def generate(self,idx,total):
+        for _ in range(total):
+            idx_cond = idx[:, -self.block_size:]
+            logits= self(idx_cond)
+            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

model/multi_head_attention.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+from .single_attention_head import AttentionHead
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, Config):
+        super().__init__()
+        self.n_heads = Config.n_heads
+        self.head_size = Config.head_size
+        
+        self.heads = nn.ModuleList([AttentionHead(Config) for _ in range(self.n_heads)])
+        
+        self.projection = nn.Linear(Config.n_embed, Config.n_embed)
+        
+        self.dropout = nn.Dropout(Config.attn_dropout)
+    
+    def forward(self,x):
+        x = torch.cat([h(x) for h in self.heads],dim=-1)
+        x = self.projection(x)
+        x = self.dropout(x)
+        return x

model/single_attention_head.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class AttentionHead(nn.Module):
+    def __init__(self, Config):
+        super().__init__()
+
+        self.block_size = Config.block_size
+        self.n_embed = Config.n_embed
+        self.head_size = Config.head_size
+        
+        self.key = nn.Linear(self.n_embed, self.head_size, bias=False)
+        self.query = nn.Linear(self.n_embed, self.head_size, bias=False)
+        
+        self.value = nn.Linear(self.n_embed, self.head_size, bias=False)
+
+        self.register_buffer(
+            'tril',
+            torch.tril(torch.ones(self.block_size,self.block_size))
+        )
+        
+        self.dropout = nn.Dropout(Config.attn_dropout)
+
+    def forward(self, x):
+
+        B,T,C = x.shape
+
+        k = self.key(x)
+        q = self.query(x)
+
+        wei = [email protected](-2,-1) * (C ** 0.5)
+        wei = wei.masked_fill(self.tril[:T,:T]==0,float('-inf'))
+        wei = F.softmax(wei, dim=-1)
+        wei = self.dropout(wei)
+        
+        v = self.value(x)
+        out = wei @ v
+        
+        return out

model/tokenizer.py

@@ -0,0 +1,16 @@
+import torch
+
+class CharacterLevelTokenizer:
+    def __init__(self,data):
+        self.data = data
+        self.vocab = sorted(list(set(self.data)))
+        self.VOCAB_SIZE = len(self.vocab)
+        
+        self.i_s = {i:s for i,s in enumerate(self.vocab)}
+        self.s_i = {s:i for i,s in self.i_s.items()}
+        
+    def encode(self,s):
+        return torch.tensor([self.s_i[c] for c in s],dtype=torch.long)
+
+    def decode(self,s):
+        return ''.join([self.i_s[i.item()] for i in s])

model/transformer_block.py

@@ -0,0 +1,18 @@
+import torch
+import torch.nn as nn
+from .multi_head_attention import MultiHeadAttention
+from .feed_forward import FeedForward
+
+class TransformerBlock(nn.Module):
+    def __init__(self, Config):
+        super().__init__()
+        self.attn = MultiHeadAttention(Config)
+        self.ff = FeedForward(Config)
+        self.ln1 = nn.LayerNorm(Config.n_embed)
+        self.ln2 = nn.LayerNorm(Config.n_embed)
+
+    def forward(self,x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.ff(self.ln2(x))
+        
+        return x

potterGPT/potterGPT.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6fe29f05742d58ac7ce20a36eea366e7022e23ef849dc5b130d185f0d36301a
+size 5733550

train.py

@@ -0,0 +1,129 @@
+from model import CharacterLevelTokenizer, Config, PotterGPT
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from pathlib import Path
+from tokenizers import Tokenizer
+import matplotlib.pyplot as plt
+
+torch.manual_seed(1357)
+with open('data/harry_potter_data', 'r', encoding='utf-8') as f:
+    data = f.read()
+
+class Dataset:
+    def __init__(self,Config, is_test=False) -> None:
+        self.tokenizer = CharacterLevelTokenizer(data)
+        self.is_test = is_test
+        self.full_data = self.tokenizer.encode(self.tokenizer.data)
+        if self.is_test:
+            self.data = self.full_data[int(0.9*len(self.full_data)):]
+        else:
+            self.data = self.full_data[:int(0.9*len(self.full_data))]
+        self.block_size = Config.block_size
+        self.batch_size = Config.batch_size
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def get_block_size(self) -> int:
+        return self.block_size
+
+    def get_vocab_size(self) -> int:
+        return self.tokenizer.VOCAB_SIZE
+
+    def get(self):
+        ix = torch.randint(len(self.data) - self.block_size, (self.batch_size,))
+        x = torch.stack([self.data[i:i+self.block_size] for i in ix])
+        y = torch.stack([self.data[i+1:i+self.block_size+1] for i in ix])
+        return x,y
+
+# tokenizer = tokenizer = Tokenizer.from_file('tokenizer/potter.json')
+tokenizer = CharacterLevelTokenizer(data)
+
+#Training
+
+train_ds = Dataset(Config)
+val_ds = Dataset(Config, is_test=True)
+
+lm = PotterGPT(Config)
+lm = lm.to(device=Config.device)
+
+optim = torch.optim.Adam(lm.parameters(), lr=Config.lr)
+
+def loss_fn(logits, targets):
+    B, T, C = logits.shape
+    logits = logits.view(B*T, C)
+    targets = targets.view(B*T)
+    loss = F.cross_entropy(logits, targets)
+    return loss
+
+def train_N_iters():
+    lm.train()
+    train_step_losses = []
+    for batch in tqdm(range(Config.train_iters)):
+        optim.zero_grad()
+        inputs, targets = train_ds.get()
+        inputs, targets = inputs.to(device=Config.device), targets.to(device=Config.device)
+        logits = lm(inputs)
+        loss = loss_fn(logits,targets)
+        loss.backward()
+        optim.step()
+        train_step_losses.append(loss.item())
+        
+        if batch%(Config.train_iters//10)==0 or batch==Config.train_iters-1:
+            print(f"batch {batch} train step loss: {loss.item()}")
+        
+        del inputs, targets, loss, logits
+        
+    return train_step_losses
+    
+@torch.no_grad()
+def valid_N_iters():
+    lm.eval()
+    val_step_losses = []
+    for batch in tqdm(range(Config.val_iters)):
+        inputs, targets = val_ds.get()
+        inputs, targets = inputs.to(device=Config.device), targets.to(device=Config.device)
+        logits = lm(inputs)
+        loss = loss_fn(logits,targets)
+        val_step_losses.append(loss.item())
+        
+        if batch%(Config.val_iters//10)==0 or batch==Config.val_iters-1:
+            print(f"batch {batch} valid step loss: {loss.item()}")
+        
+        del inputs, targets, loss, logits
+    
+    return val_step_losses
+
+def save_lm():
+    state_dict = lm.state_dict()
+    save_path = Path('./').resolve() / 'potterGPT'
+    save_path.mkdir(exist_ok=True)
+    model_path = save_path / f'potterGPT.pth'
+    torch.save(state_dict, model_path)
+
+def train_lm():
+    train_losses = train_N_iters()
+    valid_losses = valid_N_iters()
+    save_lm()
+    return train_losses, valid_losses
+
+tl, vl = train_lm()
+
+plt.plot(tl,label='train loss',color='orange')
+plt.plot(vl,label='valid loss',color='blue')
+plt.title('Potter GPT Losses')
+plt.legend()
+plt.show()
+
+generated_texts = []
+for length in [100,300,500,700,1000]:
+    generated = lm.generate(
+    torch.zeros((1,1),dtype=torch.long,device=Config.device), # initial context 0
+    total=length
+)
+    generated = tokenizer.decode(generated[0])
+    text=f'generated ({length} tokens)\n{"="*50}\n{generated}\n{"="*50}\n\n'
+    generated_texts.append(text)
+    print(text)