change structure

sentiment_clf_helper.py

@@ -4,14 +4,21 @@ from onnxruntime.quantization import quantize_dynamic,QuantType
 import transformers.convert_graph_to_onnx as onnx_convert
 from pathlib import Path
 import os
-
 import torch
-from transformers import AutoModelForSequenceClassification,AutoTokenizer
 
-chkpt='distilbert-base-uncased-finetuned-sst-2-english'
-model=AutoModelForSequenceClassification.from_pretrained(chkpt)
-tokenizer=AutoTokenizer.from_pretrained(chkpt)
 
+import yaml
+def read_yaml(file_path):
+    with open(file_path, "r") as f:
+        return yaml.safe_load(f)
+
+config = read_yaml('config.yaml')
+
+sent_chkpt=config['SENTIMENT_CLF']['sent_chkpt']
+sent_mdl_dir=config['SENTIMENT_CLF']['sent_mdl_dir']
+sent_onnx_mdl_dir=config['SENTIMENT_CLF']['sent_onnx_mdl_dir']
+sent_onnx_mdl_name=config['SENTIMENT_CLF']['sent_onnx_mdl_name']
+sent_onnx_quant_mdl_name=config['SENTIMENT_CLF']['sent_onnx_quant_mdl_name']
 
 def classify_sentiment(texts,model,tokenizer):
     """
@@ -30,7 +37,7 @@ def classify_sentiment(texts,model,tokenizer):
     return output
 
 
-def create_onnx_model_sentiment(_model, _tokenizer):
+def create_onnx_model_sentiment(_model, _tokenizer,sent_onnx_mdl_dir=sent_onnx_mdl_dir):
     """
 
     Args:
@@ -41,20 +48,21 @@ def create_onnx_model_sentiment(_model, _tokenizer):
         Creates a simple ONNX model & int8 Quantized Model in the directory "sent_clf_onnx/" if directory not present
 
     """
-    if not os.path.exists('sent_clf_onnx_dir'):
+    if not os.path.exists(sent_onnx_mdl_dir):
         try:
-            os.mkdir('sent_clf_onnx_dir')
+            os.mkdir(sent_onnx_mdl_dir)
         except:
             pass
         pipeline=transformers.pipeline("text-classification", model=_model, tokenizer=_tokenizer)
 
         onnx_convert.convert_pytorch(pipeline,
                                      opset=11,
-                                     output=Path("sent_clf_onnx_dir/sentiment_classifier_onnx.onnx"),
+                                     output=Path(f"{sent_onnx_mdl_dir}/{sent_onnx_mdl_name}"),
                                      use_external_format=False
                                      )
 
-        quantize_dynamic("sent_clf_onnx_dir/sentiment_classifier_onnx.onnx","sent_clf_onnx_dir/sentiment_classifier_onnx_quant.onnx",
+        quantize_dynamic(f"{sent_onnx_mdl_dir}/{sent_onnx_mdl_name}",
+                         f"{sent_onnx_mdl_dir}/{sent_onnx_quant_mdl_name}",
                          weight_type=QuantType.QUInt8)
     else:
         pass