File size: 3,558 Bytes
3a379e2 3ca6892 3a379e2 |
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 |
import unicodedata
import torch
import utils
import parameters
import json
from sentence_transformers import SentenceTransformer
import os
import torch.nn.functional as F
# os.environ['CUDA_VISIBLE_DEVICES'] = '0'
device = torch.device('cpu')
class Siameser:
def __init__(self, model_name=None, stadard_scope=None):
print("Load sentence embedding model (If this is the first time you run this repo, It could be take time to download sentence embedding model)")
self.threshold = 0.61
self.embedding_model = SentenceTransformer(parameters.embedding_model).to(device)
if stadard_scope == 'all':
print('Load standard address')
with open(file=parameters.NORM_ADDS_FILE_ALL_1, mode='r', encoding='utf-8') as f:
self.NORM_ADDS = json.load(fp=f)
print('Load standard address matrix')
embedding = torch.load(parameters.STD_EMBEDDING_FILE_ALL_1)
self.std_embeddings = embedding['accent_matrix'].to(device)
self.NT_std_embeddings = embedding['noaccent_matrix'].to(device)
else:
print('Load standard address')
with open(file=parameters.NORM_ADDS_FILE_HN_HCM, mode='r', encoding='utf-8') as f:
self.NORM_ADDS = json.load(fp=f)
print('Load standard address matrix')
embedding = torch.load(parameters.STD_EMBEDDING_FILE_HN_HCM)
self.std_embeddings = embedding['accent_matrix'].to(device)
self.NT_std_embeddings = embedding['noaccent_matrix'].to(device)
self.num_std_add = self.std_embeddings.shape[0]
print('Done')
def standardize(self, raw_add_):
raw_add = unicodedata.normalize('NFC', raw_add_).lower()
raw_add = utils.remove_punctuation(raw_add)
raw_add_vector = self.embedding_model.encode(raw_add, convert_to_tensor=True).to(device)
raw_add_vectors = raw_add_vector.repeat(self.num_std_add, 1)
if raw_add == utils.remove_accent(raw_add):
score = F.cosine_similarity(raw_add_vectors, self.NT_std_embeddings)
else:
score = F.cosine_similarity(raw_add_vectors, self.std_embeddings)
s, top_k = score.topk(1)
s, idx = s.tolist()[0], top_k.tolist()[0]
if s < self.threshold:
return {'Format Error': 'Xâu truyền vào không phải địa chỉ, mời nhập lại.'}
std_add = self.NORM_ADDS[str(idx)]
return utils.get_full_result(raw_add_, std_add, round(s, 4))
def get_top_k(self, raw_add_, k):
raw_add = unicodedata.normalize('NFC', raw_add_).lower()
raw_add = utils.remove_punctuation(raw_add)
raw_add_vector = self.embedding_model.encode(raw_add, convert_to_tensor=True).to(device)
raw_add_vectors = raw_add_vector.repeat(self.num_std_add, 1)
if raw_add == utils.remove_accent(raw_add):
score = F.cosine_similarity(raw_add_vectors, self.NT_std_embeddings)
else:
score = F.cosine_similarity(raw_add_vectors, self.std_embeddings)
s, top_k = score.topk(k)
s, top_k = s.tolist(), top_k.tolist()
if s[0] < self.threshold:
return {'Format Error': 'Dường như xâu truyền vào không phải địa chỉ, mời nhập lại.'}, {}
top_std_adds = []
for score, idx in zip(s, top_k):
std_add = self.NORM_ADDS[str(idx)]
top_std_adds.append(utils.get_full_result(raw_add_, std_add, round(score, 4)))
return top_std_adds[0], top_std_adds |