Settings and search calls in ElasticSearch for vector similarity search

kNN_ES_search.py

# Idea : use this to index Quora's Question Pair dataset and explore the search results
doc = {
    "settings": {
        "index.knn": True
    },
    "mappings": {
        "properties": {
            "title": {
                "type" : "text"
            },
            "title_vector": {
                "type": "knn_vector",
                "dimension": 128
            }
        }
    }
}

es.indices.create(index="questions",body=doc,ignore=400) # New index created using format specified above

def search_knn(title):
    x = np.asarray(embed([title])).tolist()[0] # embed() creates embeddings of a sentence of 128 dimensions

    script_query = {
        "knn": {
            "title_vector": {
                "vector": x,
                "k": 2 # K=2 so a maximum of 2-NN
            }
        }
    }

    response = es.search(
        index="questions",
        body={
            "size": 10, # Choose any size limit
            "query": script_query,
            "_source": {"includes": ["title"]} # Must have a title
        }
    )
    return response