Package: triton-bert
It is easy to use bert in triton inference server now.triton_bert.py
Algorithm Engineer only need to focus to write proprocess function to make his model work.model_4_triton.py
Tool for transfer huggingface pytorch.bin model into torchscript or onnx to be used in triton inference server.pgvector_triton.py
Tool for easy usage if you use postgresql/pgvector as the vector db in semantic retrievalGithubhttps://github.com/yyw794/triton-bertUSAGES:Example 0:Embedding model(Biencoder)
embedding output from model can be used directly.
so no need to override the proprocess function.fromtriton_bert.triton_bertimportTritonBertif__name__=="__main__":model=TritonBert(triton_host="30.171.160.44",model="sbert_onnx",vocab="/Users/yanyongwen712/.cache/torch/sentence_transformers/sentence-transformers_all-MiniLM-L6-v2")# batch inferencesvectors=model(["基金的收益率是多少？","我有个朋友的股票天天涨停"])# orvectors=model.encode(["基金的收益率是多少？","我有个朋友的股票天天涨停"])assertlen(vectors)==2# single inferencevector=model.encode("基金的收益率是多少？")assertvectors[0]==vectorExample 1:Embedding model(Biencoder)
Embedding need normalized, override the proprecessfromtriton_bert.triton_bertimportTritonBertimportnumpyasnpclassBiencoder(TritonBert):def__init__(self,triton_host:str,model:str,vocab:str,**kwargs):super().__init__(triton_host=triton_host,model=model,vocab=vocab,**kwargs)self.normalize_vector=Truedefproprocess(self,triton_output):ifself.normalize_vector:#if you use IP, you must normalize the vector which is the same as cosinereturn[(x/np.linalg.norm(x)).tolist()forxintriton_output[0]]returntriton_output[0].tolist()Example 2:Rank model(CrossEncoder)
user query is compared the most similar top N results with each other, and find the most similar one.fromtriton_bert.triton_bertimportTritonBertimportnumpyasnpclassCrossEncoder(TritonBert):'''rank with text similarity'''def__init__(self,triton_host:str,model:str,vocab:str,**kwargs):super().__init__(triton_host=triton_host,model=model,vocab=vocab,**kwargs)defproprocess(self,triton_output):returnnp.squeeze(triton_output[0],axis=1).tolist()def__call__(self,query,text_pairs):#change user rank input into our input pairstexts=len(text_pairs)*[query]returnself.predict(texts,text_pairs)if__name__=="__main__":model=CrossEncoder(triton_host="xx",model="xx",vocab="xx")model("小明借了小红500元",['小红借了小明500元','小明还了小红500元','小明借了小红400元'])Example 3ChitChat Intention Detection.fromtriton_bert.triton_bertimportTritonBertimporttorch.nn.functionalasFimporttorchclassChitchatIntentDetection(TritonBert):def__init__(self,**kwargs):super().__init__(**kwargs)self.label_list=["闲聊","问答","扯淡"]defproprocess(self,triton_output):logits=triton_output[0]label_ids=logits.argmax(axis=-1)logits=torch.tensor(logits)probs=F.softmax(logits,dim=1).numpy()ret=[]fori,label_idinenumerate(label_ids):prob=probs[i][label_id]iflabel_id==2andprob<0.8:label_id=0ret.append({"category":self.label_list[label_id],"confidence":float(prob)})returnretrun examplesrun triton server# for exampledockerrun-d--nametriton-server--shm-size=1g--ulimitmemlock=-1--ulimitstack=67108864--rm-p8000:8000-p8001:8001-p8002:8002-v/home/yanyongwen712/triton_models:/modelsnvcr.io/nvidia/tritonserver::22.08-py3tritonserver--model-repository=/models--model-control-mode=poll--exit-on-error=false--log-verbose1# configure triton model folderprepare model for triton serverSee the tests for more examples.Example:fromtriton_bert.model_4_tritonimportModel4TritonServerif__name__=="__main__":pretrained_model="/Users/yanyongwen712/.cache/torch/sentence_transformers/simcse-chinese-roberta-wwm-ext"model=Model4TritonServer(pretrained_model=pretrained_model)model.save_torchscript("model/simcse_model.pt")model.save_onnx("model/simcse_model.onnx")Semantic Retrieval with Postgresql/pgvectorExample 1instance=PgvectorTriton(db_user="budevadmin",db_password='eimdmPOSTGRESQL@2023',db_instance="D0PEIMDMINFO-postgresql.dbdev.paic.com.cn",db_port="3671",db_schema="pivotoperation",create_table=True,triton_host="30.171.160.44",model="bge-m3",vocab="/Users/yanyongwen712/Codes/pingan_health_rag/models/bge-m3",table_model=Sentence)# insertqas=instance.load_texts("dataset/medical_qa.jsonl")answers=[qa['answers'][0]forqainqas]instance.insert_vectors(answers)# retrievalrecalls:List[Sentence]=instance.retrieval_vectors("我喉咙有些干")print(recalls[0].sentence)