Add actual prompt instead of randomly generating text

main.py

@@ -55,8 +55,11 @@ m = BigramLanguageModel(
 ).to(DEVICE)
 
 
-def run_model(model: nn.Module, response_size: int = BLOCK_SIZE):
-    context = torch.zeros((1, 1), dtype=torch.long, device=DEVICE)
+def run_model(model: nn.Module, response_size: int = BLOCK_SIZE, query: str = ''):
+    start_ids = encode(query)
+    context = torch.tensor(start_ids, dtype=torch.long, device=DEVICE)
+    # add batch dimension. it's just 1 batch, but we still need it cuz tensors
+    context = context[None, ...]
     encoded = model.generate(
         idx=context, max_new_tokens=response_size)[0]
     return decode(encoded.tolist())
@@ -76,5 +79,5 @@ else:
     }, 'model.pth')
     print("Training complete!")
 print("Generating response...\n")
-resp = run_model(m, 256)
+resp = run_model(m, 256, 'To be or not to be, that is the question:')
 print("Response:", resp)