Upload app.py

app.py

@@ -81,10 +81,11 @@ class BigramLanguageModel(nn.Module):
         return logits, None
 
     def generate(self, idx, max_new_tokens):
+        # Ensure we respect the block size of 32
         for _ in range(max_new_tokens):
-            idx_cond = idx[:, -32:]  # Ensure context length does not exceed block size
+            idx_cond = idx[:, -32:]  # Truncate to the latest 32 tokens
             logits, _ = self(idx_cond)
-            logits = logits[:, -1, :]
+            logits = logits[:, -1, :]  # Get the logits for the last token
             probs = nn.functional.softmax(logits, dim=-1)
             idx_next = torch.multinomial(probs, num_samples=1)
             idx = torch.cat((idx, idx_next), dim=1)
@@ -105,7 +106,7 @@ model = load_model()
 chars = sorted(list(set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?-:;'\"\n")))
 stoi = {ch: i for i, ch in enumerate(chars)}
 itos = {i: ch for i, ch in enumerate(chars)}
-encode = lambda s: [stoi[c] for c in s if c in stoi]  # Ensures only known characters are encoded
+encode = lambda s: [stoi[c] for c in s if c in stoi]  # Ensure only known characters are encoded
 decode = lambda l: ''.join([itos[i] for i in l])
 
 # Function to generate text using the model
@@ -114,6 +115,10 @@ def generate_text(prompt):
         print(f"Received prompt: {prompt}")
         encoded_prompt = encode(prompt)
         
+        # Check for out-of-vocabulary indices
+        if any(idx >= 61 for idx in encoded_prompt):
+            return "Error: Input contains characters not in the model vocabulary."
+
         # Ensure the prompt length fits within the block size
         if len(encoded_prompt) > 32:
             encoded_prompt = encoded_prompt[:32]  # Truncate to fit block size
@@ -142,4 +147,4 @@ interface = gr.Interface(
 )
 
 # Launch the interface
-interface.launch(share=True)
+interface.launch()