flagembedding_bge_text_embeddings_similarity_quickstart.py

1from FlagEmbedding import FlagModel
2
3# Initialize the model
4# Models can be '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 performance hit
8
9sentences_1 = ["How is the weather today?", "What is the capital of France?"]
10sentences_2 = ["It is sunny today.", "Paris is the capital of France."]
11
12# Generate embeddings
13embeddings_1 = model.encode(sentences_1)
14embeddings_2 = model.encode(sentences_2)
15
16# Calculate similarity (inner product)
17similarity = embeddings_1 @ embeddings_2.T
18print(similarity)
19
20# For queries, it is recommended to use the encode_queries method 
21# which automatically adds the instruction
22queries = ["What is the capital of France?"]
23passages = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
24q_embeddings = model.encode_queries(queries)
25p_embeddings = model.encode(passages)
26scores = q_embeddings @ p_embeddings.T
27print(scores)