Views
No views yet
SentenceTransformer(
(0): Router(
(sub_modules): ModuleDict(
(query): Sequential(
(0): StaticEmbedding(
(embedding): EmbeddingBag(107962, 1024, mode='mean')
)
(1): Normalize()
)
(document): Sequential(
(0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertModel'})
(1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
)
)
)
)pip install -U sentence-transformers1from sentence_transformers import SentenceTransformer
2
3# Download from the 🤗 Hub
4model = SentenceTransformer("stephantulkens/NIFE-mxbai-embed-large-v1_as_router")
5# Run inference
6queries = [
7 "Which planet is known as the Red Planet?",
8]
9documents = [
10 "Venus is often called Earth's twin because of its similar size and proximity.",
11 'Mars, known for its reddish appearance, is often referred to as the Red Planet.',
12 'Saturn, famous for its rings, is sometimes mistaken for the Red Planet.',
13]
14query_embeddings = model.encode_query(queries)
15document_embeddings = model.encode_document(documents)
16print(query_embeddings.shape, document_embeddings.shape)
17# [1, 1024] [3, 1024]
18
19# Get the similarity scores for the embeddings
20similarities = model.similarity(query_embeddings, document_embeddings)
21print(similarities)
22# tensor([[0.5811, 0.7632, 0.6914]])