Delete inference.py

inference.py

@@ -1,93 +0,0 @@
-import numpy as np
-import torch
-import torch.nn.functional as F
-import tiktoken
-from dataclasses import dataclass
-
-from model import GPT
-
-
-class GPT2Inference:
-    """ To generate text sequences using a trained GPT2 model """
-
-    def __init__(self, model, token_encoder, device):
-        self.model = model
-        self.token_encoder = token_encoder
-        self.device = device
-        self.device_type = 'cuda' if device.startswith('cuda') else 'cpu'
-
-    def generate_sequences(self, prompt, num_seq=5, max_tokens=50):
-        self.model.eval()
-        tokens = self.token_encoder.encode(prompt)
-        tokens = torch.tensor(tokens, dtype=torch.long)    # (n,)   n : current sequence length
-        tokens = tokens.unsqueeze(0).repeat(num_seq, 1)    # (1,n) --> (num_seq, n)
-        gen_tokens = tokens.to(self.device)
-        # create a different rng generator so as not to impact the global rng state used for training
-        sample_rng = torch.Generator(device=self.device).manual_seed(42)
-
-        # generate new tokens one token at a time until the sequence length becomes 'max_tokens'
-        while gen_tokens.shape[-1] <= max_tokens:
-            with torch.no_grad():
-                with torch.autocast(device_type=self.device_type, dtype=torch.bfloat16):
-                    logits, loss = self.model(gen_tokens)    # (num_seq, n, vocab_size)
-                logits = logits[:, -1, :]    # (num_seq, vocab_size)
-                probs = F.softmax(logits, dim=-1)    # (num_seq, vocab_size)
-                # take top-k 50 probs
-                topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)    # (num_seq, 50), (num_seq, 50)
-                # sample a token from top-50 probabilities
-                ix = torch.multinomial(topk_probs, num_samples=1, generator=sample_rng)    # (num_seq, 1)
-                next_tok = torch.gather(topk_indices, -1, ix)    # (num_seq, 1)
-                gen_tokens = torch.cat([gen_tokens, next_tok], dim=1)
-        # decode generated tokens and print generated text
-        for i in range(num_seq):
-            tokens = gen_tokens[i, :max_tokens].tolist()
-            gen_text = self.token_encoder.decode(tokens)
-            print(f"> sample {i}: {gen_text}")
-
-
-def parse_args():
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--prompt', type=str, default="Hello, I am a language model,")
-    parser.add_argument('--num_seq', type=int, default=5)
-    parser.add_argument('--max_tokens', type=int, default=50)
-    args = parser.parse_args()
-    return args
-
-
-@dataclass
-class GPTConfig:
-    context_length: int = 1024    # max context / sequence length
-    vocab_size: int = 50257    # number of tokens: 50000 BPE merges + 256 bytes tokens + 1 <endoftext> token
-    num_layers: int = 12
-    embd_size: int = 768    # embedding dim
-    num_heads: int = 12
-
-
-def inference(args=None):
-    if args is None:
-        args = parse_args()
-
-    device = 'cpu'
-    if torch.cuda.is_available():
-        device = 'cuda'
-    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
-        device = 'mps'    # for apple macbook GPUs
-    print(f'using device: {device}')
-
-    model_path = './logs/model_95364.pt'
-    checkpoint = torch.load(model_path, weights_only=False)
-    print(f"loaded model from: {model_path}")
-    # print(checkpoint['model'].keys())
-
-    model = GPT(config=checkpoint['config'])
-    model.load_state_dict(checkpoint['model'])
-    model = model.to(device)
-    token_encoder = tiktoken.get_encoding('gpt2')
-    generator = GPT2Inference(model, token_encoder, device)
-
-    generator.generate_sequences(args.prompt, args.num_seq, args.max_tokens)
-
-
-if __name__ == '__main__':
-    inference()