update model

zeroshot_clf_helper.py

@@ -5,6 +5,19 @@ import subprocess
 import numpy as np
 import pandas as pd
 
+import yaml
+def read_yaml(file_path):
+    with open(file_path, "r") as f:
+        return yaml.safe_load(f)
+
+config = read_yaml('config.yaml')
+
+zs_chkpt=config['ZEROSHOT_CLF']['zs_chkpt']
+zs_mdl_dir=config['ZEROSHOT_CLF']['zs_mdl_dir']
+zs_onnx_mdl_dir=config['ZEROSHOT_CLF']['zs_onnx_mdl_dir']
+zs_onnx_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_mdl_name']
+zs_onnx_quant_mdl_name=config['ZEROSHOT_CLF']['zs_onnx_quant_mdl_name']
+
 
 def zero_shot_classification(premise: str, labels: str, model, tokenizer):
     try:
@@ -40,23 +53,27 @@ def zero_shot_classification(premise: str, labels: str, model, tokenizer):
 #                          labels='science, sports, museum')
 
 
-def create_onnx_model_zs(art_path='zeroshot_onnx_dir'):
+def create_onnx_model_zs(zs_onnx_mdl_dir=zs_onnx_mdl_dir):
 
     # create onnx model using
-    if not os.path.exists(art_path):
+    if not os.path.exists(zs_onnx_mdl_dir):
         try:
             subprocess.run(['python3', '-m', 'transformers.onnx',
                             '--model=valhalla/distilbart-mnli-12-1',
                             '--feature=sequence-classification',
-                            art_path])
-        except:
-            pass
+                            zs_onnx_mdl_dir])
+        except Exception as e:
+            print(e)
 
         #create quanitzed model from vanila onnx
-        quantize_dynamic(f"{art_path}/model.onnx",f"{art_path}/model_quant.onnx",weight_type=QuantType.QUInt8)
+        quantize_dynamic(f"{zs_onnx_mdl_dir}/{zs_onnx_mdl_name}",
+                         f"{zs_onnx_mdl_dir}/{zs_onnx_quant_mdl_name}",
+                         weight_type=QuantType.QUInt8)
     else:
         pass
 
+create_onnx_model_zs(zs_onnx_mdl_dir=zs_onnx_mdl_dir)
+
 def zero_shot_classification_onnx(premise,labels,_session,_tokenizer):
     try:
         labels=labels.split(',')