File size: 2,790 Bytes
046014b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import sys
sys.dont_write_bytecode = True
import unicodedata
import numpy
import sentencepiece
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]}")
|