1+ """Basic neuralnet-based language modeling"""
2+
13# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/models.lm.ipynb.
24
35# %% auto 0
4- __all__ = ['Vocab ' , 'NNLMConfig' , 'NNLM' , 'NNBigram' ]
6+ __all__ = ['ITER_MAX ' , 'NNLMConfig' , 'NNLM' , 'NNLM_L ''NNBigram' ]
57
68# %% ../../nbs/models.lm.ipynb 3
79import torch .nn as nn
810import torch
911import torch .nn .functional as F
1012from torch .nn .utils .rnn import pad_sequence
1113from torch .optim import SGD
12- from torch .utils .data import DataLoader
14+ from torch .optim .lr_scheduler import StepLR , ReduceLROnPlateau , ExponentialLR , CosineAnnealingLR
15+ from torch .optim .optimizer import Optimizer
1316
17+ from torch .utils .data import DataLoader
1418from torchtext .vocab import vocab
1519
20+ import lightning as L
21+ from lightning import Trainer
22+
1623from matplotlib import pyplot as plt
1724import pandas as pd
1825import numpy as np
26+ from tqdm import tqdm
27+
28+ from omegaconf import OmegaConf
29+ from hydra .utils import instantiate
1930
2031from typing import Dict , List , Tuple , Optional , Set
2132from collections import Counter , OrderedDict
2233from dataclasses import dataclass , asdict
2334
2435from plum import dispatch
2536
26- from ..text .datasets import CharDataset
37+ from ..text .datasets import CharDataset , Vocab
2738
28- # %% ../../nbs/models.lm.ipynb 5
29- class Vocab :
30- def __init__ (self ,
31- data :List [List [str ]], # one line per sentence. each line is a list of tokens
32- specials = ['<pad>' , '<unk>' , '<bos>' , '<eos>' ] # special characters
33- ):
34- # count individual tokens
35- c = Counter ()
36- for row in data :
37- for token in row :
38- c .update (token )
39- ordered_tuple = sorted (c .items (), key = lambda x :x [1 ], reverse = True )
40- dict = OrderedDict (ordered_tuple )
41- # leverage torchtext vocab
42- self .voc = vocab (dict , specials = specials )
43- if '<unk>' in specials :
44- self .voc .set_default_index (self .voc ['<unk>' ])
45- else :
46- self .voc .set_default_index (- 1 )
47- self ._stoi = self .voc .get_stoi ()
48- self ._itos = self .voc .get_itos ()
49-
50- @dispatch
51- def stoi (self , token :str )-> int :
52- if len (token ) > 1 and token not in '<pad>' , '<unk>' , '<bos>' , '<eos>' ]:
53- raise ValueError ("input should be a token or list of tokens" )
54- return self ._stoi [token ]
55-
56- @dispatch
57- def stoi (self , tokens :List [str ])-> List [int ]:
58- return [self ._stoi [tok ] for tok in tokens ]
59-
60- # @dispatch #TODO
61- # def stoi(self, tokens:List[List[str]])->List[List[int]]:
62- # return [self._stoi[u] for tok in tokens for ]
63- # TODO:
64- # support torch tensors
65-
66- @dispatch
67- def itos (self , index :int )-> str :
68- return self ._itos [index ]
69-
70- @dispatch
71- def itos (self , indices :List [int ])-> List [str ]:
72- return [self ._itos [index ] for index in indices ]
73-
74- def __len__ (self ):
75- return len (self .voc )
76-
77- @property
78- def vocabulary (self )-> Set :
79- return sorted (set ([k for k ,v in self ._stoi .items ()]))
80-
39+ # N_EPOCHS for training debuggging
40+ ITER_MAX = 5
8141
82- # %% ../../nbs/models.lm.ipynb 17
42+ # %% ../../nbs/models.lm.ipynb 18
8343@dataclass
8444class NNLMConfig :
8545 n_vocab :int = 30
@@ -99,15 +59,15 @@ def __init__(self,
9959 self .embedder = nn .Embedding (n_vocab , n_emb ) # (B,T)->(B,T,C)
10060 self .n_emb = n_emb
10161 self .n_context = n_context
102- # we concatenate input of n_context length * n_emb (T*C) into linear layer:
103- self .l1 = nn .Linear (n_emb * n_context , n_h )
62+ # we concatenate input of [ n_context length, n_emb] into linear layer (T*C) :
63+ self .l1 = nn .Linear (n_context * n_emb , n_h )
10464 self .l2 = nn .Linear (n_h , n_vocab )
10565
10666 def forward (self , x :torch .Tensor )-> torch .Tensor :
10767 # input: (B,T)
10868 embedding = self .embedder (x ) # ->(B,T,C)
10969 # we concatenate input of n_context length * n_emb (T*C) into linear layer:
110- h = self .l1 (embedding .view (- 1 , self .n_emb * self .n_context ))
70+ h = self .l1 (embedding .view (- 1 ,self .n_context * self .n_emb ))
11171 h = torch .tanh (h )
11272 logits = self .l2 (h )
11373 return (logits )
@@ -129,7 +89,59 @@ def sample(self, n_iterations:int=10, eos:int=3, pad:int=0, bos:int=2)->str:
12989 res .append (out )
13090 return (res )
13191
132- # %% ../../nbs/models.lm.ipynb 27
92+ # %% ../../nbs/models.lm.ipynb 36
93+ class NNLM_L (L .LightningModule ):
94+ def __init__ (
95+ self ,
96+ n_vocab :int , # vocabulary size
97+ n_emb :int , # embedding dimension
98+ n_context :int , # context size bigram/trigram, etc.
99+ n_h :int , # hidden layer size
100+ lr :float = 1e-3 , # learning rate
101+ ):
102+ super ().__init__ ()
103+ self .save_hyperparameters ()
104+ self .model = NNLM (n_vocab , n_emb , n_context , n_h )
105+ self .loss_fn = nn .CrossEntropyLoss ()
106+ self .lr = lr
107+
108+ def configure_optimizers (self ) -> Optimizer :
109+ optimizer = torch .optim .Adam (self .parameters (), lr = self .lr )
110+ return optimizer
111+
112+ def forward (self , x :torch .Tensor ) -> torch .Tensor :
113+ return self .model (x )
114+
115+ def training_step (self , batch , batch_idx ):
116+ x , y = batch
117+ y_hat = self (x )
118+ loss = self .loss_fn (y_hat , y [:, - 1 ]) # as y is shifted by one (cf. karpathy tuto)
119+ self .log ('train_loss' , loss )
120+ return loss
121+
122+ def validation_step (self , batch , batch_idx ):
123+ x , y = batch
124+ y_hat = self (x )
125+ loss = self .loss_fn (y_hat , y [:, - 1 ])
126+ self .log ('val_loss' , loss )
127+ return loss
128+
129+ def test_step (self , batch , batch_idx ):
130+ x , y = batch
131+ y_hat = self (x )
132+ loss = self .loss_fn (y_hat , y [:, - 1 ])
133+ self .log ('test_loss' , loss )
134+ return loss
135+
136+ def predict_step (self , batch , batch_idx ):
137+ x , y = batch
138+ y_hat = self (x )
139+ return y_hat
140+
141+ def sample (self , n_iterations :int = 10 , eos :int = 3 , pad :int = 0 , bos :int = 2 )-> str :
142+ return self .model .sample (n_iterations , eos , pad , bos )
143+
144+ # %% ../../nbs/models.lm.ipynb 47
133145class NNBigram (nn .Module ):
134146 def __init__ (self , vocab_size :int ) -> None :
135147 super ().__init__ ()
0 commit comments