batch dim

app.py

@@ -1,19 +1,21 @@
 import pickle
 import torch
 import gradio as gr
-from gpt import GPTLanguageModel, encode, decode
+from gpt import GPTLanguageModel
 
 
 with open('stoi_itos.pkl', 'rb') as file:
     stoi, itos = pickle.load(file)
 
+encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
+decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
 lm = GPTLanguageModel()
 lm.load_state_dict(torch.load('shakespeare_lm.pt', map_location='cpu'))
 lm.eval()
   
 
 def inference(prompt: str):
-    encoded_prompt = torch.tensor(encode(prompt), dtype=torch.long)
+    encoded_prompt = torch.tensor(encode(prompt), dtype=torch.long).unsqueeze(0) 
     output = decode(lm.generate(encoded_prompt, max_new_tokens=500)[0].tolist())
     return output
 

gpt.py

@@ -9,9 +9,7 @@ from config import gpt_config as config
 
 torch.manual_seed(1337)
 
-
-encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
-decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
+vocab_size = 65
 
 
 class Head(nn.Module):