File size: 2,330 Bytes
28e6dfa | 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 | import sys
sys.dont_write_bytecode = True
import unicodedata
import numpy
import sentencepiece
from helper import onnxSessionBuild
pathModel = "./"
embeddingTokenMax = 2048
sentencepieceEmbedding = sentencepiece.SentencePieceProcessor()
sentencepieceEmbedding.Load(f"{pathModel}tokenizer.model")
onnxSessionEmbedding = onnxSessionBuild(f"{pathModel}onnx/model.onnx")
def embedding(mode, text):
inputList = text if isinstance(text, list) else [text]
inputPrefixList = []
for a in range(len(inputList)):
if mode == "document":
inputPrefixList.append(f"title: none | text: {inputList[a]}")
else:
inputPrefixList.append(f"task: search result | query: {inputList[a]}")
tokenList = []
lengthMax = 0
for a in range(len(inputPrefixList)):
idList = sentencepieceEmbedding.EncodeAsIds(inputPrefixList[a])
if len(idList) > embeddingTokenMax - 2:
idList = idList[0:embeddingTokenMax - 2]
idList = [sentencepieceEmbedding.bos_id()] + idList + [sentencepieceEmbedding.eos_id()]
if len(idList) > lengthMax:
lengthMax = len(idList)
tokenList.append(idList)
inputIds = numpy.full((len(tokenList), lengthMax), sentencepieceEmbedding.pad_id(), dtype=numpy.int64)
attentionMask = numpy.zeros((len(tokenList), lengthMax), dtype=numpy.int64)
for a in range(len(tokenList)):
inputIds[a, 0:len(tokenList[a])] = tokenList[a]
attentionMask[a, 0:len(tokenList[a])] = 1
feedObject = {"input_ids": inputIds, "attention_mask": attentionMask}
return onnxSessionEmbedding.run(["sentence_embedding"], feedObject)[0]
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])
promptVector = embedding("query", prompt)[0]
vectorList = embedding("document", textList)
for a in range(len(textList)):
score = float(numpy.dot(promptVector, vectorList[a]) / (numpy.linalg.norm(promptVector) * numpy.linalg.norm(vectorList[a])))
print(f"{score:.6f} | {textList[a]}")
|