Repository: deepjavalibrary/djl
Pull request #3028 (2024-03-13 17:09:53)

[example] Adds document about how to trace gpt2 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


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3028?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 72.22%. Comparing base [(`0fa0db2`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/0fa0db28e481b8dbef5ec15ecfbbef4807d55b4e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) to head [(`d313389`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3028?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #3028      +/-   ##
============================================
- Coverage     72.30%   72.22%   -0.09%     
+ Complexity     7294     7290       -4     
============================================
  Files           722      722              
  Lines         32508    32508              
  Branches       3395     3395              
============================================
- Hits          23506    23479      -27     
- Misses         7382     7411      +29     
+ Partials       1620     1618       -2     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/3028?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/repository/JarRepository.java b/api/src/main/java/ai/djl/repository/JarRepository.java
index b1b443d01..1f610f99b 100644
--- a/api/src/main/java/ai/djl/repository/JarRepository.java
+++ b/api/src/main/java/ai/djl/repository/JarRepository.java
@@ -42,13 +42,16 @@ public class JarRepository extends AbstractRepository {
     private String artifactId;
     private String modelName;
     private String queryString;
+    private String originalUri;
 
     private Metadata metadata;
     private boolean resolved;
 
-    JarRepository(String name, URI uri, String fileName, String queryString) {
+    JarRepository(String name, URI uri, String fileName, URI realUri) {
         super(name, uri);
-        this.queryString = queryString;
+        this.uri = realUri;
+        queryString = uri.getRawQuery();
+        originalUri = uri.toString();
         modelName = arguments.get("model_name");
         artifactId = arguments.get("artifact_id");
         if (artifactId == null) {
@@ -123,8 +126,14 @@ public class JarRepository extends AbstractRepository {
         metadata = new Metadata.MatchAllMetadata();
         metadata.setArtifactId(artifactId);
         metadata.setArtifacts(Collections.singletonList(artifact));
-        String hash =
-                Utils.hash(queryString == null ? uri.toString() : uri.toString() + queryString);
+        String hashKey;
+        if (Boolean.parseBoolean(arguments.get("ignore_real_uri"))) {
+            hashKey = originalUri;
+        } else {
+            hashKey = queryString == null ? uri.toString() : uri.toString() + queryString;
+        }
+
+        String hash = Utils.hash(hashKey);
         MRL mrl = model(Application.UNDEFINED, DefaultModelZoo.GROUP_ID, hash);
         metadata.setRepositoryUri(mrl.toURI());
 
diff --git a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
index a730a57cb..eb5ab68a7 100644
--- a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
+++ b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
@@ -144,7 +144,6 @@ class RepositoryFactoryImpl implements RepositoryFactory {
         @Override
         public Repository newInstance(String name, URI uri) {
             String p = uri.getPath();
-            String queryString = uri.getRawQuery();
             if (p.startsWith("/")) {
                 p = p.substring(1);
             }
@@ -152,19 +151,22 @@ class RepositoryFactoryImpl implements RepositoryFactory {
             if (u == null) {
                 throw new IllegalArgumentException("Resource not found: " + uri);
             }
+
+            URI realUri;
             try {
-                uri = u.toURI();
+                // resolve real uri: jar:file:/path/my_lib.jar!/model.zip
+                realUri = u.toURI();
             } catch (URISyntaxException e) {
                 throw new IllegalArgumentException("Resource not found: " + uri, e);
             }
 
-            Path path = Paths.get(parseFilePath(uri));
+            Path path = Paths.get(parseFilePath(realUri));
             String fileName = path.toFile().getName();
             if (FilenameUtils.isArchiveFile(fileName)) {
                 fileName = FilenameUtils.getNamePart(fileName);
             }
 
-            return new JarRepository(name, uri, fileName, queryString);
+            return new JarRepository(name, uri, fileName, realUri);
         }
 
         /** {@inheritDoc} */
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
index acbaa152f..59cba679b 100644
--- a/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
@@ -59,6 +59,7 @@ public final class TextGeneration {
         SearchConfig config = new SearchConfig();
         config.setMaxSeqLength(60);
 
+        // You can use src/main/python/trace_gpt2.py to trace gpt2 model
         String url = "https://djl-misc.s3.amazonaws.com/test/models/gpt2/gpt2_pt.zip";
 
         Criteria<NDList, CausalLMOutput> criteria =
@@ -160,6 +161,20 @@ public final class TextGeneration {
         long padTokenId = 220;
         config.setPadTokenId(padTokenId);
 
+        // The model is converted optimum:
+        // https://huggingface.co/docs/optimum/main/en/exporters/onnx/usage_guides/export_a_model#exporting-a-model-using-past-keysvalues-in-the-decoder
+        /*
+         * optimum-cli export onnx --model gpt2 gpt2_onnx/
+         *
+         * from transformers import AutoTokenizer
+         * from optimum.onnxruntime import ORTModelForCausalLM
+         *
+         * tokenizer = AutoTokenizer.from_pretrained("./gpt2_onnx/")
+         * model = ORTModelForCausalLM.from_pretrained("./gpt2_onnx/")
+         * inputs = tokenizer("My name is Arthur and I live in", return_tensors="pt")
+         * gen_tokens = model.generate(**inputs)
+         * print(tokenizer.batch_decode(gen_tokens))
+         */
         String url = "https://djl-misc.s3.amazonaws.com/test/models/gpt2/gpt2_onnx.zip";
 
         Criteria<NDList, CausalLMOutput> criteria =
diff --git a/examples/src/main/python/trace_gpt2.py b/examples/src/main/python/trace_gpt2.py
new file mode 100644
index 000000000..33c3badb0
--- /dev/null
+++ b/examples/src/main/python/trace_gpt2.py
@@ -0,0 +1,73 @@
+import torch
+from transformers import GPT2LMHeadModel, GPT2Tokenizer
+
+model_name = 'gpt2-large'
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+
+# add the EOS token as PAD token to avoid warnings
+model = GPT2LMHeadModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id, torchscript=True)
+
+# %% model_inputs
+output_attentions = False
+output_hidden_states = False
+model_inputs = {}
+
+model_inputs['past_key_values'] = torch.load(
+    "../data/nested_tuple_" + model_name + ".pt")
+past_seq = model_inputs['past_key_values'][0][0].shape[-2]
+model_inputs['input_ids'] = torch.tensor([[404]])
+model_inputs['position_ids'] = torch.tensor([[past_seq]])
+# |attention_mask| = `len(past_key_values) + len(input_ids)`
+model_inputs['attention_mask'] = torch.ones(past_seq + 1, dtype=torch.int64)
+
+model_inputs['use_cache'] = True
+model_inputs['token_type_ids'] = None
+
+model_inputs['return_dict'] = False
+model_inputs['output_attentions'] = False
+model_inputs['output_hidden_states'] = False
+
+# This is a testing of text generation
+outputs = model(**model_inputs)
+
+# %% Wrapper class of GPT2LMHeadModel
+from typing import Tuple
+
+class Tracable(torch.nn.Module):
+    def __init__(self, config: dict):
+        super().__init__()
+        self.model = GPT2LMHeadModel.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id, torchscript=True)
+        self.config = {'use_cache': config.get('use_cache', True),
+                       'token_type_ids': config.get('token_type_ids', None),
+                       'return_dict': config.get('return_dict', False),
+                       'output_attentions': config.get('output_attentions', False),
+                       'output_hidden_states': config.get('output_hidden_states', True)}
+
+    def forward(self, my_input_ids, position_ids, attention_mask, past_key_values):
+        return self.model(input_ids=my_input_ids,
+                          position_ids=position_ids,
+                          attention_mask=attention_mask,
+                          past_key_values=past_key_values,
+                          **self.config)  # return_tensor = True
+
+# %% create class
+config = {}
+tracable = Tracable(config)
+input = (model_inputs['input_ids'],
+         model_inputs['position_ids'],
+         model_inputs['attention_mask'],
+         model_inputs['past_key_values'])
+
+output = tracable(*input)
+
+# %% trace
+tracable.eval()
+
+traced_model = torch.jit.trace(tracable, input)
+torch.jit.save(traced_model, "../traced_GPT2_hidden.pt")
+
+out1 = traced_model(*input)
+
+# %% load back
+loaded_model = torch.jit.load("../traced_GPT2_hidden.pt")
+out2 = loaded_model(*input)


Test patch:
diff --git a/api/src/test/java/ai/djl/repository/JarRepositoryTest.java b/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
index 4599214fa..c1370d1da 100644
--- a/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
+++ b/api/src/test/java/ai/djl/repository/JarRepositoryTest.java
@@ -45,7 +45,7 @@ public class JarRepositoryTest {
         URL[] url = {jarFile.toUri().toURL()};
         try {
             Thread.currentThread().setContextClassLoader(new URLClassLoader(url));
-            Repository repo = Repository.newInstance("test", "jar:///test.zip?hash=1");
+            Repository repo = Repository.newInstance("test", "jar:///test.zip");
             Assert.assertEquals("test", repo.getName());
             Assert.assertTrue(repo.isRemote());
 
@@ -55,6 +55,12 @@ public class JarRepositoryTest {
             Artifact artifact = repo.resolve(list.get(0), null);
             repo.prepare(artifact);
             Assert.assertEquals(1, artifact.getFiles().size());
+
+            repo = Repository.newInstance("test", "jar:///test.zip?ignore_real_uri=true");
+            list = repo.getResources();
+            artifact = repo.resolve(list.get(0), null);
+            Path p = repo.getResourceDirectory(artifact);
+            Assert.assertFalse(Files.exists(p));
         } finally {
             Thread.currentThread().setContextClassLoader(null);
         }