Update README.md

README.md

@@ -39,33 +39,34 @@ Arctic combines a 10B dense transformer model with a residual 128x3.66B MoE MLP
 total and 17B active parameters chosen using a top-2 gating. For more details about Arctic's model
 Architecture, training process, data, etc. [see our series of cookbooks](https://www.snowflake.com/en/data-cloud/arctic/cookbook/).
 
-
 ## Usage
 
-As of 4/24/2024 we are actively working with the maintainers of `transformers` to include the Arctic 
-model implementation. Until this support is released please follow these instructions to get the 
-required dependencies for using Arctic:
+Arctic is currently supported with `transformers` by leveraging the 
+[custom code feature](https://huggingface.co/docs/transformers/en/custom_models#using-a-model-with-custom-code), 
+to use this you simply need to add `trust_remote_code=True` to your AutoTokenizer and AutoModelForCausalLM calls.
+However, we recommend that you use a `transformers` version at or above 4.39:
 
 ```python
-pip install git+https://github.com/Snowflake-Labs/transformers.git@arctic
+pip install transformers>=4.39.0
 ```
 
 Arctic leverages several features from [DeepSpeed](https://github.com/microsoft/DeepSpeed), you will need to 
-install the latest version of DeepSpeed to get all of these required features:
+install the DeepSpeed 0.14.2 or higher to get all of these required features:
 
 ```python
-pip install "deepspeed>=0.14.2"
+pip install deepspeed>=0.14.2
 ```
 
-### Inference
+### Inference examples
 
 Due to the model size we recommend using a single 8xH100 instance from your
 favorite cloud provider such as: AWS [p5.48xlarge](https://aws.amazon.com/ec2/instance-types/p5/), 
 Azure [ND96isr_H100_v5](https://learn.microsoft.com/en-us/azure/virtual-machines/nd-h100-v5-series), etc.
 
 In this example we are using FP8 quantization provided by DeepSpeed in the backend, we can also use FP6 
-quantization by specifying `q_bits=6` in the `ArcticQuantizationConfig` config. The `"150GiB"` setting 
-for max_memory is required until we can get DeepSpeed's FP quantization supported natively as a [HFQuantizer](https://huggingface.co/docs/transformers/main/en/hf_quantizer#build-a-new-hfquantizer-class) which we 
+quantization by specifying `q_bits=6` in the `QuantizationConfig` config. The `"150GiB"` setting 
+for max_memory is required until we can get DeepSpeed's FP quantization supported natively as a
+[HFQuantizer](https://huggingface.co/docs/transformers/main/en/hf_quantizer#build-a-new-hfquantizer-class) which we 
 are actively working on.
 
 ```python
@@ -75,24 +76,29 @@ os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.models.arctic.configuration_arctic import ArcticQuantizationConfig
-
-tokenizer = AutoTokenizer.from_pretrained("Snowflake/snowflake-arctic-instruct")
+from deepspeed.linear.config import QuantizationConfig
 
-quant_config = ArcticQuantizationConfig(q_bits=8)
+tokenizer = AutoTokenizer.from_pretrained(
+    "Snowflake/snowflake-arctic-instruct",
+    trust_remote_code=True
+)
+quant_config = QuantizationConfig(q_bits=8)
 
 model = AutoModelForCausalLM.from_pretrained(
     "Snowflake/snowflake-arctic-instruct",
+    trust_remote_code=True,
     low_cpu_mem_usage=True,
     device_map="auto",
     ds_quantization_config=quant_config,
     max_memory={i: "150GiB" for i in range(8)},
     torch_dtype=torch.bfloat16)
 
-messages = [{"role": "user", "content": "What is 1 + 1 "}]
+
+content = "5x + 35 = 7x - 60 + 10. Solve for x"
+messages = [{"role": "user", "content": content}]
 input_ids = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cuda")
 
-outputs = model.generate(input_ids=input_ids, max_new_tokens=20)
+outputs = model.generate(input_ids=input_ids, max_new_tokens=256)
 print(tokenizer.decode(outputs[0]))
 ```