Feat: adding self.training argument to implement KV cache during validation and inference

[adityavipradas]

Resolving issue #92

• Training: do not build KV cache
• Validation: build KV cache
• Inference: build KV cache

Added self.training argument in language_model.py to populate KV cache during inference, validation and set it to [None] during training.

• This keeps the code cleaner and easy to understand compared to adding additional arguments in call functions
• but using self.training argument is not the perfect solution as it will populate KV cache during validation. Additional modifications will be needed to handle this.

[kashif]

If self.training is true during evaluation, then perhaps we are doing something wrong by missing a call to model.eval()?

[adityavipradas]

model.eval() is right where it needs to be.

What we need is:

• Training: do not build KV cache
• Validation: do not build KV cache
• Inference: build KV cache

If building KV cache during validation is fine, we can use self.training or torch.is_grad_enabled(). Please let me know and I will make the modifications accordingly. Thank you.