Update README.md

README.md

@@ -35,6 +35,7 @@ import torch
 from pathlib import Path
 from transformers import AutoModelForCausalLM
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.tokens.tokenizers.tekken import SpecialTokenPolicy
 from mistral_common.protocol.instruct.messages import UserMessage
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 from huggingface_hub import snapshot_download
@@ -68,7 +69,10 @@ def load_chat_request(message: str) -> ChatCompletionRequest:
 model_name = "prince-canuma/Ministral-8B-Instruct-2410-HF"
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 model_path = get_model_path(model_name)
+
 tokenizer = MistralTokenizer.from_file(f"{model_path}/tokenizer.json")
+tekken = tokenizer.instruct_tokenizer.tokenizer
+tekken.special_token_policy = SpecialTokenPolicy.IGNORE
 
 # Chat interaction
 user_message = "Tell me a short story about a robot learning to paint."
@@ -78,10 +82,11 @@ input_ids = torch.tensor(tokens).unsqueeze(0)
 
 # Generate response
 output = model.generate(input_ids, max_new_tokens=500, temperature=0.7, do_sample=True)
-response = tokenizer.decode(output[0][input_ids.shape[1]:])
+response = tokenizer.decode(output[0][input_ids.shape[1]:].tolist())
 
 print("User:", user_message)
 print("Model:", response)
+
 ```
 
 ## Model Details