Upload app.py

app.py

@@ -82,7 +82,7 @@ class BigramLanguageModel(nn.Module):
 
     def generate(self, idx, max_new_tokens):
         for _ in range(max_new_tokens):
-            idx_cond = idx[:, -32:]
+            idx_cond = idx[:, -32:]  # Ensure context length does not exceed block size
             logits, _ = self(idx_cond)
             logits = logits[:, -1, :]
             probs = nn.functional.softmax(logits, dim=-1)
@@ -112,11 +112,19 @@ decode = lambda l: ''.join([itos[i] for i in l])
 def generate_text(prompt):
     try:
         print(f"Received prompt: {prompt}")
-        context = torch.tensor([encode(prompt)], dtype=torch.long)
+        encoded_prompt = encode(prompt)
+        
+        # Ensure the prompt length fits within the block size
+        if len(encoded_prompt) > 32:
+            encoded_prompt = encoded_prompt[:32]  # Truncate to fit block size
+        
+        context = torch.tensor([encoded_prompt], dtype=torch.long)
         print(f"Encoded prompt: {context}")
+        
         with torch.no_grad():
             generated = model.generate(context, max_new_tokens=250)  # Adjust as needed
             print(f"Generated tensor: {generated}")
+        
         result = decode(generated[0].tolist())
         print(f"Decoded result: {result}")
         return result