Upload app.py

app.py

@@ -81,7 +81,6 @@ 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:]  # Truncate to the latest 32 tokens
             logits, _ = self(idx_cond)
@@ -103,10 +102,11 @@ def load_model():
 model = load_model()
 
 # Define a comprehensive character set based on training data
-chars = sorted(list(set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789 .,!?-:;'\"\n")))
+# Convert all input to lowercase if the model is trained on lowercase data
+chars = sorted(list(set("abcdefghijklmnopqrstuvwxyz0123456789 .,!?-:;'\"\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]  # Ensure only known characters are encoded
+encode = lambda s: [stoi.get(c, stoi.get(c.lower(), -1)) for c in s if c in stoi or c.lower() in stoi]  # Handles both cases
 decode = lambda l: ''.join([itos[i] for i in l])
 
 # Function to generate text using the model
@@ -116,7 +116,7 @@ def generate_text(prompt):
         encoded_prompt = encode(prompt)
         
         # Check for out-of-vocabulary indices
-        if any(idx >= 61 for idx in encoded_prompt):
+        if any(idx == -1 for idx in encoded_prompt):
             return "Error: Input contains characters not in the model vocabulary."
 
         # Ensure the prompt length fits within the block size