Update README.md

README.md

@@ -19,7 +19,7 @@ The following `bitsandbytes` quantization config was used during training:
 
 - PEFT 0.4.0.dev0
 
-
+### Import model
 ```python
 import torch
 from peft import PeftModel, PeftConfig
@@ -32,7 +32,9 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
+```
 
+### Model Inference
 
 ```python
 input_text = "participant safety -->: "
@@ -42,3 +44,4 @@ with torch.cuda.amp.autocast():
     output_tokens = model.generate(**batch, max_new_tokens=50)
 
 print('\n\n', tokenizer.decode(output_tokens[0], skip_special_tokens=True))
+```