RAG知识检索

Halcom

RAG知识检索：
BAAI/bge-small-zh-v1.5 · HF Mirror

1. from transformers import AutoTokenizer, AutoModel
   
2. import torch
   
3. # Sentences we want sentence embeddings for
   
4. sentences = ["样例数据-1", "样例数据-2"]
   
5. 
   
6. # Load model from HuggingFace Hub
   
7. tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-large-zh-v1.5')
   
8. model = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5')
   
9. model.eval()
   
10. 
    
11. # Tokenize sentences
    
12. encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
    
13. # for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
    
14. # encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
    
15. 
    
16. # Compute token embeddings
    
17. with torch.no_grad():
    
18.     model_output = model(**encoded_input)
    
19.     # Perform pooling. In this case, cls pooling.
    
20.     sentence_embeddings = model_output[0][:, 0]
    
21. # normalize embeddings
    
22. sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
    
23. print("Sentence embeddings:", sentence_embeddings)

复制代码

**encoded_input是三个tensor组合的。可以展开：类似这种model_output = model(tensor(encoded_input["input_ids"]), tensor(encoded_input["mask"]), tensor(encoded_input["ids"]))

Halcom

encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')也可以写为：
encoded_input = tokenizer(sentences)

下载训练好的模型，直接config设置模型参数，
config = AutoConfig.from_pretrained(model_dir)
model = AutoModel.from_config(config)

模型输入等效：
model_output = model( input_ids=torch.tensor(encoded_input['input_ids']), attention_mask=torch.tensor(encoded_input['attention_mask']), token_type_ids=torch.tensor(encoded_input['token_type_ids']) )

Halcom

模型转化，
torch.save(model.state_dict(), 'rag_'+ str(512) + '.pt')
model.load_state_dict(torch.load('rag_'+ str(512) + '.pt', map_location=torch.device('cpu') ))
model.eval()

example1 = torch.randint(low=0, high=1000, size=(1, 512), dtype=torch.long).type(torch.LongTensor)
example2 = torch.ones(1, 512).type(torch.LongTensor)
example3 = torch.zeros(1, 512).type(torch.LongTensor)

pt_model = torch.jit.trace(model, (example1,example2,example3), strict=False )
pt_model.save( 'ragX_'+ str(512) + '.pt' )

Halcom

1. # -*- coding: utf-8 -*-
   
2. """
   
3. Spyder Editor
   
4. This is a temporary script file.
   
5. """
   
6. from transformers import AutoTokenizer, AutoModel, AutoConfig
   
7. import torch
   
8. # Sentences we want sentence embeddings for
   
9. sentences = ["RAG模型已初始化"]
   
10. 
    
11. model_dir = r'D:/2-LearningCode/913_LM\RAG_WITIAI/bge-small-zh-1.5'
    
12. # Load model from HuggingFace Hub
    
13. tokenizer = AutoTokenizer.from_pretrained(model_dir)
    
14. # model = AutoModel.from_pretrained('BAAI/bge-large-zh-v1.5')
    
15. config = AutoConfig.from_pretrained(model_dir)
    
16. model = AutoModel.from_config(config)
    
17. model.eval()
    
18. 
    
19. # torch.save(model.state_dict(), 'rag_'+ str(512) + '.pt')
    
20. # model.load_state_dict(torch.load('rag_'+ str(512) + '.pt', map_location=torch.device('cpu') ))
    
21. # model.eval()
    
22. 
    
23. example1 = torch.randint(low=0, high=1000, size=(1, 512), dtype=torch.long).type(torch.LongTensor)
    
24. example2 = torch.ones(1, 512).type(torch.LongTensor)
    
25. example3 = torch.zeros(1, 512).type(torch.LongTensor)
    
26. 
    
27. pt_model = torch.jit.trace(model, (example1,example2,example3), strict=False )
    
28. pt_model.save( 'ragX_'+ str(512) + '.pt' )
    
29. 
    
30. # 第二步，直接load，C++中是这样加载
    
31. # model = torch.jit.load('ragX_'+ str(512) + '.pt')
    
32. # model.eval()
    
33. 
    
34. # Tokenize sentences
    
35. # encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
    
36. # for s2p(short query to long passage) retrieval task, add an instruction to query (not add instruction for passages)
    
37. # encoded_input = tokenizer([instruction + q for q in queries], padding=True, truncation=True, return_tensors='pt')
    
38. encoded_input = tokenizer(sentences)
    
39. 
    
40. # Compute token embeddings
    
41. with torch.no_grad():
    
42.     # model_output = model(**encoded_input)
    
43.     model_output = model( input_ids=torch.tensor(encoded_input['input_ids']),
    
44.                          attention_mask=torch.tensor(encoded_input['attention_mask']),
    
45.                          token_type_ids=torch.tensor(encoded_input['token_type_ids']) )
    
46.     # Perform pooling. In this case, cls pooling.
    
47.     sentence_embeddings = model_output[0][:, 0]
    
48. # normalize embeddings
    
49. sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
    
50. print("Sentence embeddings:", sentence_embeddings)
    
51. 
    
52. 
    

复制代码

Halcom

python中测试，可用：

1. from transformers import AutoTokenizer, AutoModel, AutoConfig
   
2. import torch
   
3. 
   
4. sentences = ["RAG模型已初始化"]
   
5. 
   
6. model_dir = r'D:/2-LearningCode/913_LM\RAG_WITIAI/bge-small-zh-1.5'
   
7. tokenizer = AutoTokenizer.from_pretrained(model_dir)
   
8. config = AutoConfig.from_pretrained(model_dir)
   
9. model = AutoModel.from_config(config)
   
10. 
    
11. model.load_state_dict(torch.load('rag_'+ str(512) + '.pt', map_location=torch.device('cpu') ))
    
12. model.eval()
    
13. 
    
14. encoded_input = tokenizer(sentences)
    
15. 
    
16. # Compute token embeddings
    
17. with torch.no_grad():
    
18.     # model_output = model(**encoded_input)
    
19.     model_output = model( input_ids=torch.tensor(encoded_input['input_ids']),
    
20.                          attention_mask=torch.tensor(encoded_input['attention_mask']),
    
21.                          token_type_ids=torch.tensor(encoded_input['token_type_ids']) )
    
22.     # Perform pooling. In this case, cls pooling.
    
23.     sentence_embeddings = model_output[0][:, 0]
    
24. # normalize embeddings
    
25. sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1)
    
26. print("Sentence embeddings:", sentence_embeddings)

复制代码