

import torch
from transformers import AutoTokenizer, AutoModel


if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")
# device = torch.device("cpu")


# 加载预训练模型和分词器
# from transformers import BertTokenizer,BertModel
# tokenizer = BertTokenizer.from_pretrained('bert-large-chinese')
# model = BertModel.from_pretrained('bert-large-chinese')


model_path = "bert-base-chinese"
model_path = "THUDM/glm-large-chinese"
model_path = "E:/huggingface/hub/models--THUDM--glm-large-chinese/snapshots/230f54e413fab4bc8f29bd3508aab301d757ef3e"
model_path = "/data/python/models--THUDM--glm-large-chinese/snapshots/230f54e413fab4bc8f29bd3508aab301d757ef3e/"
tokenizer = AutoTokenizer.from_pretrained(model_path,trust_remote_code=True)
model = AutoModel.from_pretrained(model_path,trust_remote_code=True)
model = model.to(device)


def getVectorOfSentence(sentence):
    inputs = tokenizer.batch_encode_plus([sentence], return_tensors='pt', padding=True, truncation=True)["input_ids"]
    inputs = inputs.to(device)

    input_shape = inputs.size()
    position_ids = torch.arange(0, input_shape[-1], dtype=torch.long, device=device)
    block_position_ids = torch.zeros(input_shape[-1], dtype=torch.long, device=device)
    position_ids = torch.stack((position_ids, block_position_ids), dim=0).unsqueeze(0)
    with torch.no_grad():
        outputs = model(inputs,position_ids=position_ids)
    last_hidden_state = outputs.last_hidden_states
    print(last_hidden_state.shape)
    # cls_embeddings = last_hidden_state[:,-1,:]
    cls_embeddings = torch.mean(last_hidden_state,dim=1)
    return cls_embeddings

def get_cosine_similarity(v1,v2):
    return torch.dot(v1,v2)/(torch.norm(v1)*torch.norm(v2)+1e-8)


def normalize_vector(v):
    return v/torch.sqrt(torch.sum(v*v,dim=1))

def cosine_similarity(x1, x2, dim=1, eps=1e-8):
    """
    计算两个张量在指定维度上的余弦相似度。
    """
    w1 = torch.norm(x1, p=2, dim=dim, keepdim=True)
    w2 = torch.norm(x2, p=2, dim=dim, keepdim=True)
    dot = (x1 * x2).sum(dim=dim, keepdim=True)
    cos = dot / (w1 * w2 + eps)

    return cos.squeeze(dim=dim)

def get_normalized_vector(sentence):
    vector = normalize_vector(getVectorOfSentence(sentence)).tolist()[0]
    return vector


def test():
    import time
    a = time.time()
    v1 = getVectorOfSentence("123456？")
    print("get vector takes %.3f seconds" % (time.time() - a))
    a = time.time()
    v2 = getVectorOfSentence("12346")
    print("get vector takes %.3f seconds" % (time.time() - a))
    print(v1.shape)
    a = time.time()
    print(cosine_similarity(v1,v2))
    print("cosine_similarity takes %.3f seconds" % (time.time() - a))

if __name__ == '__main__':
    test()