Update README.md

README.md

@@ -47,26 +47,35 @@ You can then start chatting with the model, *e.g.* prompt it with something like
 To use this model within the `transformers` library, install the latest release with `pip install --upgrade transformers` and run, for instance:
 
 ```py
-from transformers import MistralForCausalLM
-from transformers import AutoTokenizer
+from transformers import pipeline
+import torch
 
-tokenizer = AutoTokenizer.from_pretrained('mistralai/mathstral-7B-v0.1')
+checkpoint = "mistralai/mathstral-7B-v0.1"
+pipe = pipeline("text-generation", checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
 
-prompt = "What are the roots of unity?"
-tokenized_prompts = tokenizer(prompt, return_tensors="pt") 
+prompt = [{"role": "user", "content": "What are the roots of unity?"}]
+out = pipe(prompt, max_new_tokens = 512)
 
-model = MistralForCausalLM.from_pretrained('mistralai/mathstral-7B-v0.1')
-generation = model.generate(**tokenized_prompts, max_new_tokens=512)
-print(tokenizer.decode(generation[0]))
->>> """<s>What are the roots of unity?
+print(out[0]['generated_text'][-1])
+>>> "{'role': 'assistant', 'content': ' The roots of unity are the complex numbers that satisfy the equation $z^n = 1$, where $n$ is a positive integer. These roots are evenly spaced around the unit circle in the complex plane, and they have a variety of interesting properties and applications in mathematics and physics.'}"
+```
+
+You can also manually tokenize the input and generate text from the model, rather than using the higher-level pipeline:
+
+```py
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-The roots of unity are the solutions to the equation $z^n = 1$, where $n$ is a positive integer.
-These roots are complex numbers and they form a regular $n$-gon in the complex plane.
+checkpoint = "mistralai/mathstral-7B-v0.1"
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
 
-For example, the roots of unity for $n=1$ are just $1$,
-and for $n=2$ they are $1$ and $-1$. For $n=3$, they are $1$, $\\frac{-1+\\sqrt{3}i}{2}$, and $\\frac{-1-\\sqrt{3}i}{2}$.
+prompt = [{"role": "user", "content": "What are the roots of unity?"}]
+tokenized_prompt = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, return_dict=True, return_tensors="pt").to(model.device)
 
-The roots of unity have many interesting properties and they are used in many areas of mathematics, including number theory, algebra, and geometry.</s>"""
+out = model.generate(**tokenized_prompt, max_new_tokens=512)
+tokenizer.decode(out[0])
+>>> '<s>[INST] What are the roots of unity?[/INST] The roots of unity are the complex numbers that satisfy the equation $z^n = 1$, where $n$ is a positive integer. These roots are evenly spaced around the unit circle in the complex plane, and they have a variety of interesting properties and applications in mathematics and physics.</s>'
 ```
 
 ## Evaluation