Update README.md

README.md

@@ -26,6 +26,7 @@ Despite the powerful capabilities of this model, users should be aware of its li
 How to use
 Here is an example of how to use this model:
 
+```python
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import time
 import torch
@@ -33,9 +34,7 @@ import torch
 class Chatbot:
     def __init__(self, model_name):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, 
-                                                          load_in_4bit=True,
-                                                          torch_dtype=torch.bfloat16)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name, load_in_4bit=True, torch_dtype=torch.bfloat16)
         if self.tokenizer.pad_token_id is None:
             self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
             
@@ -45,9 +44,9 @@ class Chatbot:
             inputs = {name: tensor.to('cuda') for name, tensor in inputs.items()}
         start_time = time.time()
         tokens = self.model.generate(input_ids=inputs['input_ids'], 
-                                    attention_mask=inputs['attention_mask'],
-                                    pad_token_id=self.tokenizer.pad_token_id,
-                                    max_new_tokens=400)
+                                     attention_mask=inputs['attention_mask'],
+                                     pad_token_id=self.tokenizer.pad_token_id,
+                                     max_new_tokens=400)
         end_time = time.time()
         output_tokens = tokens[0][inputs['input_ids'].shape[-1]:]
         output = self.tokenizer.decode(output_tokens, skip_special_tokens=True)