Views
No views yet
pip install sparsembed1from sparsembed import model, retrieve
2from transformers import AutoModelForMaskedLM, AutoTokenizer
3
4device = "cuda" # cpu
5
6batch_size = 10
7
8# List documents to index:
9documents = [
10 {'id': 0,
11 'title': 'Paris',
12 'url': 'https://en.wikipedia.org/wiki/Paris',
13 'text': 'Paris is the capital and most populous city of France.'},
14 {'id': 1,
15 'title': 'Paris',
16 'url': 'https://en.wikipedia.org/wiki/Paris',
17 'text': "Since the 17th century, Paris has been one of Europe's major centres of science, and arts."},
18 {'id': 2,
19 'title': 'Paris',
20 'url': 'https://en.wikipedia.org/wiki/Paris',
21 'text': 'The City of Paris is the centre and seat of government of the region and province of Île-de-France.'
22}]
23
24model = model.Splade(
25 model=AutoModelForMaskedLM.from_pretrained("raphaelsty/splade-max").to(device),
26 tokenizer=AutoTokenizer.from_pretrained("raphaelsty/splade-max"),
27 device=device
28)
29
30retriever = retrieve.SpladeRetriever(
31 key="id", # Key identifier of each document.
32 on=["title", "text"], # Fields to search.
33 model=model # Splade retriever.
34)
35
36retriever = retriever.add(
37 documents=documents,
38 batch_size=batch_size,
39 k_tokens=256, # Number of activated tokens.
40)
41
42retriever(
43 ["paris", "Toulouse"], # Queries
44 k_tokens=20, # Maximum number of activated tokens.
45 k=100, # Number of documents to retrieve.
46 batch_size=batch_size
47)