实战入门FAISS搜索BERT模型 -- 知识铺

1

docker run -p 8888:8888 myencyclopedia/faiss-demo bash -c 'jupyter notebook --allow-root --port 8888 --NotebookApp.token= --ip 0.0.0.0'

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

import faiss 


index = faiss.IndexFlatL2(d)

 
index.add(xb)                 
print(index.ntotal) 



k = 4

D, I = index.search(xq, k)     


print(I[:5])

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

import faiss
nlist = 100
k = 4


quantizer = faiss.IndexFlatL2(d)


index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)


index.train(xb) 
index.add(xb)                
index.nprobe = 10  
D, I = index.search(xq, k)
print(I[-5:])   

1
2
3
4
5
6
7
8
9
10
11
12
13
14

import faiss

nlist = 100

m = 8 

k = 4 
quantizer = faiss.IndexFlatL2(d)    
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) 
index.train(xb)
index.add(xb)
index.nprobe = 10                
D, I = index.search(xq, k)  
print(I[-5:])

1
2
3
4
5
6
7
8
9
10

def load_dataset(part='test') -> List[str]:
    ds = datasets.load_dataset("ag_news")
    list_str = [r['text'] for r in ds[part]]
    return list_str
    
list_str = load_dataset(part='train')
print(f'{len(list_str)}')
for s in list_str[:3]:
    print(s)
    print('\n')

1
2
3
4
5
6
7
8
9

120000
Wall St. Bears Claw Back Into the Black (Reuters) Reuters - Short-sellers, Wall Street's dwindling\band of ultra-cynics, are seeing green again.


Carlyle Looks Toward Commercial Aerospace (Reuters) Reuters - Private investment firm Carlyle Group,\which has a reputation for making well-timed and occasionally\controversial plays in the defense industry, has quietly placed\its bets on another part of the market.


Oil and Economy Cloud Stocks' Outlook (Reuters) Reuters - Soaring crude prices plus worries\about the economy and the outlook for earnings are expected to\hang over the stock market next week during the depth of the\summer doldrums.

1
2
3
4

index = load_index('news_train.index')
list_id = query(model, index, 'how to make money')
for id in list_id:
    print(list_str[id])

1
2
3
4
5
6
7
8
9

Profit From That Traffic Ticket Got a traffic ticket? Can't beat 'em? Join 'em by investing in the company that processes those tickets.

Answers in the Margins By just looking at operating margins, investors can find profitable industry leaders.

Types of Investors: Which Are You? Learn a little about yourself, and it may improve your performance.

Target Can Aim High Target can maintain its discount image while offering pricier services and merchandise.

Finance moves Ford into the black US carmaker Ford Motor returns to profit, as the money it makes from lending to customers outweighs losses from selling vehicles.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

def train_flat(index_name, id_list, embedding_list, num_clusters):
    import numpy as np
    import faiss

    dim = 768
    m = 16
    
    embeddings = np.asarray(embedding_list)
    
    quantiser = faiss.IndexFlatIP(dim)
    index = faiss.IndexIVFFlat(quantiser, dim, num_clusters, faiss.METRIC_INNER_PRODUCT)
    index.train(embeddings)  
    
    ids = np.arange(0, len(id_list))
    ids = np.asarray(ids.astype('int64'))
    
    index.add_with_ids(embeddings, ids)
    print(index.is_trained) 
    print("Total Number of Embeddings in the index", index.ntotal)
    faiss.write_index(index, index_name)

1
2
3
4
5
6
7

def query(model, index, query_str: str) -> List[int]:
    topk = 5
    q_embed = model.encode([query_str])
    D, I = index.search(q_embed, topk)
    print(D)
    print(I)
    return I[0].tolist()