fix tokenizer in example

TorchMoE · Jan 18, 2025 · 7aade14 · 7aade14
1 parent 218f152
commit 7aade14
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/interface_example.py b/examples/interface_example.py
@@ -11,7 +11,7 @@
 import multiprocessing as mp
 from transformers import AutoTokenizer, TextStreamer, LlamaTokenizerFast
 from moe_infinity import MoE
-from moe_infinity.models.arctic import ArcticTokenizer
+from moe_infinity.models.modeling_arctic import ArcticTokenizer
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--model_name_or_path", type=str, required=True)
@@ -27,7 +27,7 @@
 elif "arctic" in args.model_name_or_path.lower():
     tokenizer = ArcticTokenizer.from_pretrained(args.model_name_or_path)
 else:
-    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True)
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, trust_remote_code=True, use_fast=False)
 streamer = TextStreamer(tokenizer)
 
 dataset_name = "tasksource/bigbench"