Package: abottle
abottletrition/tensorrt/onnxruntim/pytorch python server wrapperput your model intoa bottlethen you get a working server and more.DemoimportnumpyasnpfromtransformersimportAutoTokenizerclassMiniLM:def__init__(self):self.tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")defpredict(self,X):encode_dict=self.tokenizer(X,padding="max_length",max_length=128,truncation=True)input_ids=np.array(encode_dict["input_ids"],dtype=np.int32)attention_mask=np.array(encode_dict["attention_mask"],dtype=np.int32)outputs=self.model.infer({"input_ids":input_ids,"attention_mask":attention_mask},["y"])returnoutputs['y']#you can write config in class or provide it as a yaml file or yaml stringclassConfig:classTritonModel:name="minilm"version="2"you can write a class like this, and then starts with abottleabottlemain.MiniLMwith default, abottle will run as server, and server at 0.0.0.0:8081curllocalhost:8081/predictabottle will inject an attribute namedmodelinto your class, and you don't need to care what that model runtime is.
it can be Pytorch with CuDNN8 or an optimized TensorRT plan, it depends on the config you giveself.model.infer({"input1":input1_tensor,"input2":input2_tensor},['output_1'])config with shellabottlemain.MiniLM--config"""TritonModel:triton_url: localhostname: minilmversion: 2"""config with fileabottlemain.MiniLM--config<configyamlfilepath>importnumpyasnpimportpandasaspdfromtransformersimportAutoTokenizerfromtypingimportListclassMiniLM:def__init__(self):self.tokenizer=AutoTokenizer.from_pretrained("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")defcosine(self,a:List[List[float]],b:List[List[float]])->float:a,b=np.array(a),np.array(b)# |A|sqrt_sqare_A=np.tile(np.sqrt(np.sum(np.square(a),axis=1)).reshape((a.shape[0],1)),(1,a.shape[0]),)# |B|sqrt_sqare_B=np.tile(np.sqrt(np.sum(np.square(b.T),axis=0)).reshape((1,b.shape[0])),(b.shape[0],1),)# cosine similarityscore_matrix=np.divide(np.dot(a,b.T),sqrt_sqare_A*sqrt_sqare_B)returnscore_matrixdefpredict(self,X:List[str])->List[List[float]]:encode_dict=self.tokenizer(X,padding="max_length",max_length=128,truncation=True)input_ids=np.array(encode_dict["input_ids"],dtype=np.int32)attention_mask=np.array(encode_dict["attention_mask"],dtype=np.int32)outputs=self.model.infer({"input_ids":input_ids,"attention_mask":attention_mask},["y"])returnoutputs["y"]defevaluate(self,file_path:str,batch_size:int)->float:test_data=pd.read_csv(file_path,sep=", ",names=["query","label"])query,label=test_data["query"].tolist(),test_data["label"].tolist()assertlen(query)==len(label)query_embedding,label_embedding=[],[]foriinrange(0,len(query),batch_size):query_embedding+=self.predict(query[i:min(i+batch_size,len(query))])label_embedding+=self.predict(label[i:min(i+batch_size,len(label))])assertlen(query_embedding)==len(label_embedding)# 分数矩阵score_matrix=self.cosine(query_embedding,label_embedding)# 算法性能raw_result=np.argmax(score_matrix,axis=0)==np.array([iforiinrange(score_matrix.shape[0])])unique,counts=np.unique(a,return_counts=True)top_1_accuracy=counts[unique.tolist().index(True)]/np.sum(counts)returntop_1_accuracydef evaluate can be used as a tester like belowabottlemain.MiniLM--astesterfile_path='test.csv',batch_size=100the arguments you defined in theevaluatefunction can be set in CLI args with format xxx=xxxyou can use different wrapper for your model, including:abottle.ONNXModelabottle.TensorRTModelabottle.TritonModelabottle.PytorchModelif you want to add more wrappers you can just implement abottle.BaseModelabottlemain.MiniLM--asserver--wrapperabottle.TritonModelConfigsabottle.ONNXModelONNXModel:ort_file:'theortfilepath'abottle.TensorRTModelTensorRTModel:trt_file:'TensorRTplanfilepath'abottle.TritonModelTritonModel:name:"yourmodel'snameontritonserver"version:"yourmodel'sversionontritonserver"triton_url:"tritonserver'shostwithoutschema,itmeanshttp://xxxisinvalid"abottle.PytorchModel(not fully implemented)PytorchModel:model:'pytrochimportablename'Motivationas a DL model creator, you don't need to focus on how to serve or test the performance of a model on a target platform or how to optimize your model and don't lose accuracy, just find a bottle and put your logic code into it, the DL engineer people can do those things for you, all you need to do is export your model to a onnx file, and write logic code like above examples.Featurewe will build this bottle as strong as possible, make this bottle become a standardization interface of the MLOps cycles, you can see more and more scenarios like optimization, graph fusing, performance test, deployment, data gathering, etc using this bottle.