Spaces:
Build error
Build error
| import os | |
| import torch | |
| import pickle | |
| class Model: | |
| def __init__(self, words): | |
| self.words = words | |
| # Create the vocabulary and the necessary mapprings | |
| self.chars = sorted(set(''.join(words))) | |
| self.stoi = {s:i+1 for i,s in enumerate(self.chars)} # mapping from string to index | |
| self.stoi['.'] = 0 | |
| self.itos = {i:s for s,i in self.stoi.items()} # mapping from index to string | |
| self.vocab_size = len(self.chars) + 1 # 26 characters + 1 special token | |
| def train(self): | |
| # Create the bigram probability matrix | |
| N = torch.zeros((self.vocab_size, self.vocab_size), dtype=torch.int32) | |
| for w in self.words: | |
| chs = ['.'] + list(w) + ['.'] | |
| for ch1, ch2 in zip(chs, chs[1:]): | |
| ix1, ix2 = self.stoi[ch1], self.stoi[ch2] | |
| N[ix1, ix2] += 1 | |
| self.P = (N+1).float() | |
| self.P /= self.P.sum(1, keepdims=True) | |
| loss = self._evaluate_likelihood() | |
| print(f"Average negative log likelihood: {loss}") | |
| def generate(self, n_names): | |
| names = [] | |
| # Generate a series of n_names tokens | |
| for i in range(n_names): | |
| out = [] | |
| ix = 0 | |
| while True: | |
| p = self.P[ix].float() | |
| ix = torch.multinomial(p, num_samples=1, replacement=True).item() | |
| if ix == 0: | |
| break | |
| else: | |
| out.append(self.itos[ix]) | |
| names.append(''.join(out)) | |
| return ', '.join(names) | |
| def save_model(self, file_path='.'): | |
| with open(os.path.join(file_path, 'model.pkl'), 'wb') as f: | |
| pickle.dump(self, f) | |
| def _evaluate_likelihood(self): | |
| log_likelihood = 0.0 | |
| n = 0 | |
| for w in self.words: | |
| chs = ['.'] + list(w) + ['.'] | |
| for ch1, ch2 in zip(chs, chs[1:]): | |
| ix1 = self.stoi[ch1] | |
| ix2 = self.stoi[ch2] | |
| prob = self.P[ix1, ix2] | |
| logprob = torch.log(prob) | |
| log_likelihood += logprob | |
| n += 1 | |
| return -log_likelihood/n | |
| if __name__ == '__main__': | |
| # Read in the words | |
| words = open('names.txt', 'r').read().splitlines() | |
| # Instantiate the model class | |
| model = Model(words) | |
| # Train the model | |
| model.train() | |
| model.save_model() | |