| import sys |
| import unicodedata |
| import numpy |
| import sentencepiece |
|
|
| sys.dont_write_bytecode = True |
|
|
| |
| from helper import onnxSessionBuild |
|
|
| pathModel = "./" |
|
|
| rerankerBatchLength = 8 |
| rerankerTokenMax = 512 |
| rerankerPromptTokenMax = 128 |
| rerankerBosId = 0 |
| rerankerPadId = 1 |
| rerankerEosId = 2 |
| rerankerUnkId = 3 |
| rerankerOffset = 1 |
|
|
| sentencepieceReranker = sentencepiece.SentencePieceProcessor() |
| sentencepieceReranker.Load(f"{pathModel}sentencepiece.bpe.model") |
|
|
| onnxSessionReranker = onnxSessionBuild(f"{pathModel}onnx/model.onnx") |
|
|
| def rerankTokenize(text): |
| resultList = [] |
|
|
| for spmId in sentencepieceReranker.encode(text, out_type=int): |
| if spmId == 0: |
| resultList.append(rerankerUnkId) |
| else: |
| resultList.append(spmId + rerankerOffset) |
|
|
| return resultList |
|
|
| def rerank(prompt, textList): |
| scoreList = [] |
|
|
| promptIdList = rerankTokenize(prompt) |
|
|
| if len(promptIdList) > rerankerPromptTokenMax: |
| promptIdList = promptIdList[0:rerankerPromptTokenMax] |
|
|
| for a in range(0, len(textList), rerankerBatchLength): |
| batchList = textList[a:a + rerankerBatchLength] |
|
|
| sequenceList = [] |
| lengthMax = 0 |
|
|
| for b in range(len(batchList)): |
| textIdList = rerankTokenize(batchList[b]) |
|
|
| lengthText = rerankerTokenMax - len(promptIdList) - 4 |
|
|
| if len(textIdList) > lengthText: |
| textIdList = textIdList[0:lengthText] |
|
|
| idList = [rerankerBosId] + promptIdList + [rerankerEosId, rerankerEosId] + textIdList + [rerankerEosId] |
|
|
| if len(idList) > lengthMax: |
| lengthMax = len(idList) |
|
|
| sequenceList.append(idList) |
|
|
| inputIds = numpy.full((len(sequenceList), lengthMax), rerankerPadId, dtype=numpy.int64) |
| attentionMask = numpy.zeros((len(sequenceList), lengthMax), dtype=numpy.int64) |
|
|
| for b in range(len(sequenceList)): |
| inputIds[b, 0:len(sequenceList[b])] = sequenceList[b] |
| attentionMask[b, 0:len(sequenceList[b])] = 1 |
|
|
| feedObject = {"input_ids": inputIds, "attention_mask": attentionMask} |
|
|
| logits = onnxSessionReranker.run(["logits"], feedObject)[0] |
|
|
| for b in range(len(logits)): |
| scoreList.append(float(1.0 / (1.0 + numpy.exp(-logits[b][0])))) |
|
|
| return scoreList |
|
|
| prompt = unicodedata.normalize("NFKC", "what is panda?") |
|
|
| textList = [ |
| "The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear, is a bear species endemic to China.", |
| "hi", |
| "パンダはクマ科の哺乳類で、中国の固有種である。" |
| ] |
|
|
| for a in range(len(textList)): |
| textList[a] = unicodedata.normalize("NFKC", textList[a]) |
|
|
| scoreList = rerank(prompt, textList) |
|
|
| for a in range(len(textList)): |
| print(f"{scoreList[a]:.6f} | {textList[a]}") |
|
|
