Merge pull request Snowflake-Labs#11 from Snowflake-Labs/trust_remote_code

jeffra · web-flow · commit 686f1bac2236 · 2024-04-26T14:54:37.000-07:00
Update inference tutorials to use trust_remote_code instead of transformers fork
diff --git a/inference/README.md b/inference/README.md
@@ -12,7 +12,6 @@ now you will need to use our forks.
 
 ```bash
 deepspeed>=0.14.2
-git+git://github.com/Snowflake-Labs/transformers.git@arctic
 git+git://github.com/Snowflake-Labs/vllm.git@arctic
 huggingface_hub[hf_transfer]
 ```
@@ -38,15 +37,19 @@ os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers.models.arctic.configuration_arctic import ArcticQuantizationConfig
+from deepspeed.linear.config import QuantizationConfig
 
-tokenizer = AutoTokenizer.from_pretrained("Snowflake/snowflake-arctic-instruct")
+tokenizer = AutoTokenizer.from_pretrained(
+    "Snowflake/snowflake-arctic-instruct",
+    trust_remote_code=True
+)
 
-quant_config = ArcticQuantizationConfig(q_bits=8)
+quant_config = QuantizationConfig(q_bits=8)
 
 model = AutoModelForCausalLM.from_pretrained(
     "Snowflake/snowflake-arctic-instruct",
     low_cpu_mem_usage=True,
+    trust_remote_code=True,
     device_map="auto",
     ds_quantization_config=quant_config,
     max_memory={i: "150GiB" for i in range(8)},
diff --git a/inference/requirements.txt b/inference/requirements.txt
@@ -1,4 +1,2 @@
 deepspeed>=0.14.2
-git+git://github.com/Snowflake-Labs/transformers.git@arctic
-git+git://github.com/Snowflake-Labs/vllm.git@arctic
-huggingface_hub[hf_transfer]
+huggingface_hub[hf_transfer]
