flagembedding_bge_dense_sentence_embeddings_similarity.py

1from FlagEmbedding import FlagModel
2
3# Initialize the model
4# You can choose different model sizes (e.g., 'BAAI/bge-large-en-v1.5', 'BAAI/bge-base-en-v1.5', etc.)
5model = FlagModel('BAAI/bge-large-en-v1.5', 
6                  query_instruction_for_retrieval="Represent this sentence for searching relevant passages:",
7                  use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight loss in precision
8
9sentences_1 = ["How do I bake a chocolate cake?", "What is the capital of France?"]
10sentences_2 = ["Instructions for making a chocolate dessert.", "Paris is the capital city of France."]
11
12# Generate embeddings
13embeddings_1 = model.encode(sentences_1)
14embeddings_2 = model.encode(sentences_2)
15
16# Compute similarity (inner product)
17similarity = embeddings_1 @ embeddings_2.T
18print(similarity)
19
20# For retrieval tasks, use encode_queries for the queries and encode for the documents
21queries = ["how to cook a cake"]
22passages = ["You can bake a cake in the oven.", "The weather is nice today."]
23q_embeddings = model.encode_queries(queries)
24p_embeddings = model.encode(passages)
25scores = q_embeddings @ p_embeddings.T
26print(scores)