[tokenizer] Refactor djl_convert python code

extensions/tokenizers/.gitignore

@@ -1,3 +1,6 @@
 Cargo.lock
 /tokenizers
 /jnilib
+model/
+tmp/
+models.json

extensions/tokenizers/src/main/python/.gitignore

@@ -1,4 +1,2 @@
 __pycache__
-model/
-tmp/
-models.json
+*.egg-info/

extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class FillMaskConverter(HuggingfaceConverter):

extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py

@@ -17,14 +17,14 @@
 from argparse import Namespace
 
 import onnx
-import safetensors_convert
+from djl_converter.safetensors_convert import convert_file
 import torch
 from huggingface_hub import hf_hub_download, HfApi
 from transformers import pipeline, AutoTokenizer, AutoConfig
 
-from metadata import HuggingfaceMetadata
-from shasum import sha1_sum
-from zip_utils import zip_dir
+from djl_converter.metadata import HuggingfaceMetadata
+from djl_converter.shasum import sha1_sum
+from djl_converter.zip_utils import zip_dir
 
 
 class PipelineHolder(object):
@@ -139,7 +139,7 @@ def save_rust_model(self, model_info, args: Namespace, temp_dir: str):
  elif has_pt_file:
  file = hf_hub_download(repo_id=model_id,
  filename="pytorch_model.bin")
- safetensors_convert.convert_file(file, target)
+ convert_file(file, target)
  else:
  return False, f"No model file found for: {model_id}", -1
 

extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py

@@ -19,6 +19,12 @@
 from huggingface_hub import HfApi
 from huggingface_hub import hf_hub_download
 from huggingface_hub.hf_api import ModelInfo
+from djl_converter.fill_mask_converter import FillMaskConverter
+from djl_converter.metadata import get_lang_tags
+from djl_converter.question_answering_converter import QuestionAnsweringConverter
+from djl_converter.sentence_similarity_converter import SentenceSimilarityConverter
+from djl_converter.text_classification_converter import TextClassificationConverter
+from djl_converter.token_classification_converter import TokenClassificationConverter
 
 ARCHITECTURES_2_TASK = {
  "ForQuestionAnswering": "question-answering",
@@ -27,19 +33,13 @@
  "ForMultipleChoice": "text-classification",
  "ForMaskedLM": "fill-mask",
 }
-LANGUAGES = HfApi().get_model_tags()["language"]
-
-
-def get_lang_tags(model_info):
- tags = {}
- for tag in model_info.tags:
- if tag in LANGUAGES:
- tags[tag] = "true"
-
- if not tags:
- tags["en"] = "true"
-
- return tags
+SUPPORTED_TASKS = {
+ "fill-mask": FillMaskConverter(),
+ "question-answering": QuestionAnsweringConverter(),
+ "sentence-similarity": SentenceSimilarityConverter(),
+ "text-classification": TextClassificationConverter(),
+ "token-classification": TokenClassificationConverter(),
+}
 
 
 class HuggingfaceModels:

extensions/tokenizers/src/main/python/djl_converter/metadata.py

@@ -11,7 +11,22 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import json
-from huggingface_models import get_lang_tags
+
+from huggingface_hub import HfApi
+
+LANGUAGES = HfApi().get_model_tags()["language"]
+
+
+def get_lang_tags(model_info):
+ tags = {}
+ for tag in model_info.tags:
+ if tag in LANGUAGES:
+ tags[tag] = "true"
+
+ if not tags:
+ tags["en"] = "true"
+
+ return tags
 
 
 class HuggingfaceMetadata:

extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py

@@ -11,25 +11,13 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import logging
-import os.path
+import os
 import shutil
 import sys
 
-from arg_parser import converter_args
-from fill_mask_converter import FillMaskConverter
-from huggingface_models import HuggingfaceModels
-from question_answering_converter import QuestionAnsweringConverter
-from sentence_similarity_converter import SentenceSimilarityConverter
-from text_classification_converter import TextClassificationConverter
-from token_classification_converter import TokenClassificationConverter
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
 
-SUPPORTED_TASK = {
- "fill-mask": FillMaskConverter(),
- "question-answering": QuestionAnsweringConverter(),
- "sentence-similarity": SentenceSimilarityConverter(),
- "text-classification": TextClassificationConverter(),
- "token-classification": TokenClassificationConverter(),
-}
+from djl_converter.arg_parser import converter_args
 
 
 def main():
@@ -38,6 +26,8 @@ def main():
  level=logging.INFO)
  args = converter_args()
 
+ from djl_converter.huggingface_models import HuggingfaceModels, SUPPORTED_TASKS
+
  huggingface_models = HuggingfaceModels(args.output_dir)
  temp_dir = f"{args.output_dir}/tmp"
 
@@ -48,7 +38,7 @@ def main():
  for model in models:
  task = model["task"]
  model_info = model["model_info"]
- converter = SUPPORTED_TASK[task]
+ converter = SUPPORTED_TASKS[task]
 
  try:
  result, reason, size = converter.save_model(

extensions/tokenizers/src/main/python/djl_converter/question_answering_converter.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class QuestionAnsweringConverter(HuggingfaceConverter):

extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py

@@ -19,7 +19,7 @@
 import torch
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
-from huggingface_converter import HuggingfaceConverter, PipelineHolder
+from djl_converter.huggingface_converter import HuggingfaceConverter, PipelineHolder
 from huggingface_hub import hf_hub_download
 
 

extensions/tokenizers/src/main/python/djl_converter/text_classification_converter.py

@@ -15,7 +15,7 @@
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class TextClassificationConverter(HuggingfaceConverter):

extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py

@@ -14,7 +14,7 @@
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class TokenClassificationConverter(HuggingfaceConverter):