Update README.md

add missing code-block langs

README.md

@@ -115,7 +115,7 @@ model = model.to(device)
 ```
 
 Load dataset
-```
+```python
 from datasets import load_dataset
 
 # Load dataset
@@ -127,7 +127,7 @@ validation_dataset = squad_dataset['test']
 ```
 
 Preprocessing: tokenize inputs and labels for faster training cycles, i.e. no need for tokenization during training anymore
-```
+```python
 def preprocess_batch(batch, tokenizer, max_input_length=512, max_output_length=128):
     contexts = batch['context']
     answers = batch['question']