Upload README.md

README.md

@@ -24,19 +24,20 @@ model = FlaxMistralForCausalLM.from_pretrained("rdyro/Mistral-7B-Instruct-v0.1",
 
 tokenizer = AutoTokenizer.from_pretrained("rdyro/Mistral-7B-Instruct-v0.1")
 
-torch_model_id = "mistralai/Mistral-7B-Instruct-v0.1"
-torch_model = AutoModelForCausalLM.from_pretrained(
-    torch_model_id, device_map="cpu", torch_dtype=torch.float32
-)
-torch_tokenizer = AutoTokenizer.from_pretrained(torch_model_id)
+messages = [{"role": "user", "content": "what's your name?"}]
+input_jax = tokenizer.apply_chat_template(messages, return_tensors="jax")
 out_jax = model(input_jax)
 ```
 
 We can compare the outputs to the original PyTorch version.
 
 ```python
-messages = [{"role": "user", "content": "what's your name?"}]
-input_jax = tokenizer.apply_chat_template(messages, return_tensors="jax")
+torch_model_id = "mistralai/Mistral-7B-Instruct-v0.1"
+torch_model = AutoModelForCausalLM.from_pretrained(
+    torch_model_id, device_map="cpu", torch_dtype=torch.float32
+)
+torch_tokenizer = AutoTokenizer.from_pretrained(torch_model_id)
+
 input_pt = torch_tokenizer.apply_chat_template(messages, return_tensors="pt")
 
 with torch.no_grad():
@@ -46,6 +47,7 @@ err = jnp.linalg.norm(jnp.array(out_pt.logits) - out_jax.logits) / jnp.linalg.no
     jnp.array(out_pt.logits)
 )
 print(f"Error is numerical precision level: {err:.4e}")
+# prints: Error is numerical precision level: 1.0205e-06
 ```
 
 <p align="center">