Update README.md

README.md

@@ -23,7 +23,7 @@ inference: false
 To use this model, install the required packages:
 
 ```bash
-pip install -U mistral-common transformers torch
+pip install -U transformers
 ```
 
 ## Usage Example
@@ -31,62 +31,25 @@ pip install -U mistral-common transformers torch
 Here's a Python script demonstrating how to use the model for chat completion:
 
 ```python
-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
-
-def get_model_path(path_or_hf_repo: str, revision: str = None) -> Path:
-    """Ensures the model is available locally, downloading if necessary."""
-    model_path = Path(path_or_hf_repo)
-    if not model_path.exists():
-        model_path = Path(
-            snapshot_download(
-                repo_id=path_or_hf_repo,
-                revision=revision,
-                allow_patterns=[
-                    "*.json",
-                    "*.safetensors",
-                    "*.py",
-                    "tokenizer.model",
-                    "*.tiktoken",
-                    "*.txt",
-                ],
-                resume_download=True,
-            )
-        )
-    return model_path
-
-def load_chat_request(message: str) -> ChatCompletionRequest:
-    """Creates a chat completion request with a user message."""
-    return ChatCompletionRequest(messages=[UserMessage(content=message)])
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Model setup
 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
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Chat interaction
-user_message = "Tell me a short story about a robot learning to paint."
-completion_request = load_chat_request(user_message)
-tokens = tokenizer.encode_chat_completion(completion_request).tokens
-input_ids = torch.tensor(tokens).unsqueeze(0)
+prompt = "Tell me a short story about a robot learning to paint."
+messages = [{"role": "user", "content": prompt}]
+text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(text, return_tensors="pt").to(model.device)
 
 # 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]:].tolist())
+output = model.generate(**input_ids, max_new_tokens=500, temperature=0.7, do_sample=True)
+response = tokenizer.decode(output[0][input_ids.input_ids.shape[1]:])
 
-print("User:", user_message)
+print("User:", prompt)
 print("Model:", response)
-
 ```
 
 ## Model Details