flagembedding

import os from FlagEmbedding import FlagAutoModel # You can replace 'BAAI/bge-base-en-v1.5' with other BGE models like 'BAAI/bge-m3' # Consider setting query_instruction_for_retrieval for optimal performance in retrieval tasks. # Use use_fp16=True for faster inference on compatible hardware. # Example for embedding queries and passages model = FlagAutoModel.from_pretrained( 'BAAI/bge-base-en-v1.5', query_instruction_for_retrieval="Represent this sentence for searching relevant passages:", use_fp16=True ) queries = ["What is FlagEmbedding?", "How to use embedding models?"] passages = [ "FlagEmbedding maps text to low-dimensional dense vectors for tasks like retrieval.", "Embedding models can be used to generate vector representations of text.", "The BGE models are state-of-the-art embedding models." ] # Encode queries and passages query_embeddings = model.encode_queries(queries) passage_embeddings = model.encode_corpus(passages) print(f"Query embeddings shape: {query_embeddings.shape}") print(f"Passage embeddings shape: {passage_embeddings.shape}") # Compute similarity scores (e.g., dot product) scores = query_embeddings @ passage_embeddings.T print("Similarity scores:") print(scores) # Example for reranking using FlagReranker from FlagEmbedding import FlagReranker # Replace with a reranker model like 'BAAI/bge-reranker-base' reranker = FlagReranker('BAAI/bge-reranker-base', use_fp16=True) query_passage_pairs = [ ['What is AI?', 'Artificial intelligence (AI) is intelligence demonstrated by machines.'], ['What is AI?', 'The quick brown fox jumps over the lazy dog.'] ] ranks = reranker.compute_score(query_passage_pairs) print("Reranker scores:") print(ranks)