Views
No views yet
ColBERT(
(0): Transformer({'max_seq_length': 32, 'do_lower_case': False, 'architecture': 'T5EncoderModel'})
(1): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': False})
)pip install -U pylate1from pylate import indexes, models, retrieve
2
3# Step 1: Load the ColBERT model
4model = models.ColBERT(
5 model_name_or_path="xtr_pylate",
6)
7
8# Step 2: Initialize the PLAID index
9index = indexes.PLAID(
10 index_folder="pylate-index",
11 index_name="index",
12 override=True, # This overwrites the existing index if any
13)
14
15# Step 3: Encode the documents
16documents_ids = ["1", "2", "3"]
17documents = ["document 1 text", "document 2 text", "document 3 text"]
18
19documents_embeddings = model.encode(
20 documents,
21 batch_size=32,
22 is_query=False, # Ensure that it is set to False to indicate that these are documents, not queries
23 show_progress_bar=True,
24)
25
26# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
27index.add_documents(
28 documents_ids=documents_ids,
29 documents_embeddings=documents_embeddings,
30)1# To load an index, simply instantiate it with the correct folder/name and without overriding it
2index = indexes.PLAID(
3 index_folder="pylate-index",
4 index_name="index",
5)1# Step 1: Initialize the ColBERT retriever
2retriever = retrieve.ColBERT(index=index)
3
4# Step 2: Encode the queries
5queries_embeddings = model.encode(
6 ["query for document 3", "query for document 1"],
7 batch_size=32,
8 is_query=True, # # Ensure that it is set to False to indicate that these are queries
9 show_progress_bar=True,
10)
11
12# Step 3: Retrieve top-k documents
13scores = retriever.retrieve(
14 queries_embeddings=queries_embeddings,
15 k=10, # Retrieve the top 10 matches for each query
16)1from pylate import rank, models
2
3queries = [
4 "query A",
5 "query B",
6]
7
8documents = [
9 ["document A", "document B"],
10 ["document 1", "document C", "document B"],
11]
12
13documents_ids = [
14 [1, 2],
15 [1, 3, 2],
16]
17
18model = models.ColBERT(
19 model_name_or_path="xtr_pylate",
20)
21
22queries_embeddings = model.encode(
23 queries,
24 is_query=True,
25)
26
27documents_embeddings = model.encode(
28 documents,
29 is_query=False,
30)
31
32reranked_documents = rank.rerank(
33 documents_ids=documents_ids,
34 queries_embeddings=queries_embeddings,
35 documents_embeddings=documents_embeddings,
36)