| from .exec_backends.trt_loader import TrtModel, encode as encode_trt | |
| from transformers import AutoTokenizer | |
| import math | |
| tokenizer_en = AutoTokenizer.from_pretrained("tensorRT/models/paraphrase-mpnet-base-v2") | |
| model_en = TrtModel("tensorRT/models/paraphrase-mpnet-base-v2.engine") | |
| tokenizer_multilingual = AutoTokenizer.from_pretrained("tensorRT/models/paraphrase-multilingual-MiniLM-L12-v2") | |
| model_multilingual= TrtModel("tensorRT/models/paraphrase-multilingual-MiniLM-L12-v2.engine") | |
| def encode(sentences, lang, batch_size = 8): | |
| if batch_size >=8: | |
| batch_size = 8 | |
| all_embs = [] | |
| NUM_BATCH = math.ceil(len(sentences) / batch_size) | |
| for j in range(NUM_BATCH): | |
| lst_sen = sentences[j*batch_size: j*batch_size + batch_size] | |
| if lang == 'en': | |
| # print(lst_sen) | |
| embs = encode_trt(lst_sen, tokenizer=tokenizer_en, trt_model= model_en, use_token_type_ids=False) | |
| else: | |
| # print(lst_sen) | |
| embs = encode_trt(lst_sen, tokenizer=tokenizer_multilingual, trt_model= model_multilingual, use_token_type_ids=False) | |
| all_embs.extend(embs) | |
| return all_embs | |