Repository: deepjavalibrary/djl
Pull request #3150 (2024-05-02 22:01:07)

[tokenizer] Allows use HF_TOKEN to access gated model
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/extensions/tokenizers/rust/src/lib.rs b/extensions/tokenizers/rust/src/lib.rs
index 6e427f683..af3fec7cd 100644
--- a/extensions/tokenizers/rust/src/lib.rs
+++ b/extensions/tokenizers/rust/src/lib.rs
@@ -48,13 +48,18 @@ pub extern "system" fn Java_ai_djl_huggingface_tokenizers_jni_TokenizersLibrary_
     mut env: JNIEnv<'local>,
     _: JObject,
     input: JString,
+    hf_token: JString,
 ) -> jlong {
     let identifier: String = env
         .get_string(&input)
         .expect("Couldn't get java string!")
         .into();
 
-    let parameters = FromPretrainedParameters::default();
+    let mut parameters = FromPretrainedParameters::default();
+    if !hf_token.is_null() {
+        let hf_token: String = env.get_string(&hf_token).unwrap().into();
+        parameters.auth_token = Some(hf_token);
+    }
     let tokenizer = Tokenizer::from_pretrained(identifier, Some(parameters));
 
     match tokenizer {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
index 8d32ccb4b..c146f5d46 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
@@ -112,8 +112,12 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
     public static HuggingFaceTokenizer newInstance(String identifier, Map<String, String> options) {
         Ec2Utils.callHome("Huggingface");
         LibUtils.checkStatus();
+        String autoToken = Utils.getEnvOrSystemProperty("HF_TOKEN");
+        if (options != null) {
+            autoToken = options.getOrDefault("hf_token", autoToken);
+        }
 
-        long handle = TokenizersLibrary.LIB.createTokenizer(identifier);
+        long handle = TokenizersLibrary.LIB.createTokenizer(identifier, autoToken);
         return new HuggingFaceTokenizer(handle, options);
     }
 
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
index 63a2f96fd..8acfa7d04 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/jni/TokenizersLibrary.java
@@ -20,7 +20,7 @@ public final class TokenizersLibrary {
 
     private TokenizersLibrary() {}
 
-    public native long createTokenizer(String identifier);
+    public native long createTokenizer(String identifier, String authToken);
 
     public native long createTokenizerFromString(String json);
 


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
index 1b0883dfb..ead688539 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizerTest.java
@@ -14,6 +14,7 @@
 package ai.djl.huggingface.tokenizers;
 
 import ai.djl.huggingface.tokenizers.jni.CharSpan;
+import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.DownloadUtils;
 import ai.djl.util.PairList;
 import ai.djl.util.Utils;
@@ -541,4 +542,17 @@ public class HuggingFaceTokenizerTest {
             Assert.assertEquals(tokens, expected);
         }
     }
+
+    @Test
+    public void testAuthToken() {
+        TestRequirements.notOffline();
+        System.setProperty("HF_TOKEN", "test_token");
+        try {
+            HuggingFaceTokenizer.newInstance("mistralai/Mixtral-8x7B-v0.1");
+        } catch (RuntimeException ignore) {
+            // access denied
+        } finally {
+            System.clearProperty("HF_TOKEN");
+        }
+    }
 }