app update

app.py

@@ -155,6 +155,9 @@ class BigramLanguageModel(nn.Module):
         self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
         self.ln_f = nn.LayerNorm(n_embd)
         self.lm_head = nn.Linear(n_embd, self.vocab_size)
+        self.encode = lambda s: [stoi[c] for c in s] # encoder: take a string, output a list of integers
+        self.decode = lambda l: ''.join([itos[i] for i in l]) # decoder: take a list of integers, output a string
+
         
     def forward(self, idx, targets=None):
         B, T = idx.shape
@@ -236,7 +239,7 @@ wikipedia_model.eval()  # Set the model to evaluation mode
 
 def generate_shakespeare_outputs(prompt=None, max_new_tokens=2000):
   if prompt:
-    context = torch.tensor(encode(prompt), dtype=torch.long, device=device).view(1, -1)
+    context = torch.tensor(shakespeare_model.encode(prompt), dtype=torch.long, device=device).view(1, -1)
   else:
     context = torch.zeros((1, 1), dtype=torch.long, device=device)
   text_output = decode(shakespeare_model.generate(context, max_new_tokens=max_new_tokens)[0].tolist())
@@ -245,7 +248,7 @@ def generate_shakespeare_outputs(prompt=None, max_new_tokens=2000):
 
 def generate_wikipedia_outputs(prompt=None, max_new_tokens=2000):
   if prompt:
-    context = torch.tensor(encode(prompt), dtype=torch.long, device=device).view(1, -1)
+    context = torch.tensor(wikipedia_model.encode(prompt), dtype=torch.long, device=device).view(1, -1)
   else:
     context = torch.zeros((1, 1), dtype=torch.long, device=device)
   text_output = decode(wikipedia_model.generate(context, max_new_tokens=max_new_tokens)[0].tolist())