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' )