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"]))