Update README.md

README.md

@@ -17,11 +17,16 @@ It has been trained using [TRL](https://github.com/huggingface/trl).
 ## Quick start
 
 ```python
-from transformers import pipeline
+import torch
+from transformers import pipeline, AutoTokenizer
 
 question = "Write a Python function that takes a list of numbers and returns the list sorted in ascending order without using the built-in `sorted()` function."
-generator = pipeline("text-generation", model="mariasandu/python-coding-assistant-v2", device="cuda")
-output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
+generator = pipeline("text-generation", model="mariasandu/python-coding-assistant-v2", device="cpu")
+
+tokenizer = AutoTokenizer.from_pretrained("mariasandu/python-coding-assistant-v2")
+tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}[INST] {{ message['content'] }} [/INST]{% else %}{{ message['content'] }}{% endif %}{% endfor %}"
+formatted_chat = tokenizer.apply_chat_template(chat, tokenize=True, return_dict=True, continue_final_message=True)
+output = generator(question)
 print(output["generated_text"])
 ```