Views
No views yet
1import os
2from transformers import (
3 LlamaForCausalLM,
4 LlamaTokenizer, AutoConfig,
5)
6import torch
7import torch.nn.functional as F
8import numpy as np
9
10
11class ZNVEmbeddingModel(torch.nn.Module):
12 def __init__(self, model_name_or_path):
13 super(ZNVEmbeddingModel, self).__init__()
14 self.prompt_prefix = "阅读下文，然后答题\n"
15 self.prompt_suffixes = ["\n1.一个字总结上文的意思是:",
16 "\n2.上文深层次的意思是:"]
17 self.hidden_size = 4096
18 self.model_name_or_path = model_name_or_path
19 self.linear_suffixes = torch.nn.ModuleList(
20 [torch.nn.Linear(self.hidden_size, self.hidden_size//len(self.prompt_suffixes))
21 for _ in range(len(self.prompt_suffixes))])
22 self.tokenizer, self.llama = self.load_llama()
23
24 self.tanh = torch.nn.Tanh()
25 self.suffixes_ids = []
26 self.suffixes_ids_len = []
27 self.suffixes_len = 0
28 for suffix in self.prompt_suffixes:
29 ids = self.tokenizer(suffix, return_tensors="pt")["input_ids"].tolist()[0]
30 self.suffixes_ids += ids
31 self.suffixes_ids_len.append(len(ids))
32 self.suffixes_len += len(ids)
33
34 self.suffixes_ones = torch.ones(self.suffixes_len)
35 self.suffixes_ids = torch.tensor(self.suffixes_ids)
36
37 linear_file = os.path.join(model_name_or_path, "linears")
38 load_layers = torch.load(linear_file)
39 model_state = self.state_dict()
40 model_state.update(load_layers)
41 self.load_state_dict(model_state, strict=False)
42
43 def load_llama(self):
44 llm_path = os.path.join(self.model_name_or_path)
45 config = AutoConfig.from_pretrained(llm_path)
46 tokenizer = LlamaTokenizer.from_pretrained(self.model_name_or_path)
47 tokenizer.padding_side = "left"
48 model = LlamaForCausalLM.from_pretrained(
49 llm_path,
50 config=config,
51 low_cpu_mem_usage=True
52 )
53 model.config.use_cache = False
54 return tokenizer, model
55
56 def forward(self, sentences):
57 prompts_embeddings = []
58 sentences = [self.prompt_prefix + s for s in sentences]
59 inputs = self.tokenizer(sentences, max_length=256, padding=True, truncation=True,
60 return_tensors='pt')
61 attention_mask = inputs["attention_mask"]
62 input_ids = inputs["input_ids"]
63 batch_size = len(sentences)
64 suffixes_ones = self.suffixes_ones.unsqueeze(0)
65 suffixes_ones = suffixes_ones.repeat(batch_size, 1)
66 device = next(self.parameters()).device
67 attention_mask = torch.cat([attention_mask, suffixes_ones], dim=-1).to(device)
68
69 suffixes_ids = self.suffixes_ids.unsqueeze(0)
70 suffixes_ids = suffixes_ids.repeat(batch_size, 1)
71 input_ids = torch.cat([input_ids, suffixes_ids], dim=-1).to(device)
72 last_hidden_state = self.llama.base_model.base_model(attention_mask=attention_mask, input_ids=input_ids).last_hidden_state
73 index = -1
74 for i in range(len(self.suffixes_ids_len)):
75 embedding = last_hidden_state[:, index, :]
76 embedding = self.linear_suffixes[i](embedding)
77 prompts_embeddings.append(embedding)
78 index -= self.suffixes_ids_len[-i-1]
79
80 output_embedding = torch.cat(prompts_embeddings, dim=-1)
81 output_embedding = self.tanh(output_embedding)
82 output_embedding = F.normalize(output_embedding, p=2, dim=1)
83 return output_embedding
84
85 def encode(self, sentences, batch_size=10, **kwargs):
86 size = len(sentences)
87 embeddings = None
88 handled = 0
89 while handled < size:
90 tokens = sentences[handled:handled + batch_size]
91 output_embeddings = self.forward(tokens)
92 result = output_embeddings.cpu().numpy()
93 handled += result.shape[0]
94 if embeddings is not None:
95 embeddings = np.concatenate((embeddings, result), axis=0)
96 else:
97 embeddings = result
98 return embeddings1znv_model = ZNVEmbeddingModel("your_model_path")
2znv_model.eval()
3with torch.no_grad():
4 output = znv_model(["请问你的电话号码是多少？","可以告诉我你的手机号吗？"])
5 cos_sim = F.cosine_similarity(output[0],output[1],dim=0)
6 print(cos_sim)