Spaces:
Runtime error
Runtime error
| import numpy as np | |
| import transformers | |
| from onnxruntime.quantization import quantize_dynamic,QuantType | |
| import transformers.convert_graph_to_onnx as onnx_convert | |
| from pathlib import Path | |
| import os | |
| import torch | |
| 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): | |
| """ | |
| user will pass texts separated by comma | |
| """ | |
| try: | |
| texts=texts.split(',') | |
| except: | |
| pass | |
| input = tokenizer(texts, padding=True, truncation=True, | |
| return_tensors="pt") | |
| logits = model(**input)['logits'].softmax(dim=1) | |
| logits = torch.argmax(logits, dim=1) | |
| output = ['Positive' if i == 1 else 'Negative' for i in logits] | |
| return output | |
| def create_onnx_model_sentiment(_model, _tokenizer,sent_onnx_mdl_dir=sent_onnx_mdl_dir): | |
| """ | |
| Args: | |
| _model: model checkpoint with AutoModelForSequenceClassification | |
| _tokenizer: model checkpoint with AutoTokenizer | |
| Returns: | |
| Creates a simple ONNX model & int8 Quantized Model in the directory "sent_clf_onnx/" if directory not present | |
| """ | |
| if not os.path.exists(sent_onnx_mdl_dir): | |
| try: | |
| 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(f"{sent_onnx_mdl_dir}/{sent_onnx_mdl_name}"), | |
| use_external_format=False | |
| ) | |
| # 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 | |
| def classify_sentiment_onnx(texts, _session, _tokenizer): | |
| """ | |
| Args: | |
| texts: input texts from user | |
| _session: pass ONNX runtime session | |
| _tokenizer: Relevant Tokenizer e.g. AutoTokenizer.from_pretrained("same checkpoint as the model") | |
| Returns: | |
| list of Positve and Negative texts | |
| """ | |
| try: | |
| texts=texts.split(',') | |
| except: | |
| pass | |
| _inputs = _tokenizer(texts, padding=True, truncation=True, | |
| return_tensors="np") | |
| input_feed={ | |
| "input_ids":np.array(_inputs['input_ids']), | |
| "attention_mask":np.array((_inputs['attention_mask'])) | |
| } | |
| output = _session.run(input_feed=input_feed, output_names=['output_0'])[0] | |
| output=np.argmax(output,axis=1) | |
| output = ['Positive' if i == 1 else 'Negative' for i in output] | |
| return output | |