txtai Similarity really slow with ElasticSearch

I’ve noticed when running ElasticSearch and txtai.pipeline for Similarity, the search (ranksearch) is very slow. When trying to search for 1 item, it can take upto 10 seconds.

The code I’m using is:

from txtai.pipeline import Similarity
from elasticsearch import Elasticsearch, helpers

# Connect to ES instance
es = Elasticsearch(hosts=["http://localhost:9200"], timeout=60, retry_on_timeout=True)

def ranksearch(query, limit):
  results = [text for _, text in search(query, limit * 10)]
  return [(score, results[x]) for x, score in similarity(query, results)][:limit]

def search(query, limit):
  query = {
      "size": limit,
      "query": {
          "query_string": {"query": query}
      }
  }

  results = []
  for result in es.search(index="articles", body=query)["hits"]["hits"]:
    source = result["_source"]
    results.append((min(result["_score"], 18) / 18, source["title"]))
  return results

similarity = Similarity("valhalla/distilbart-mnli-12-3")

limit = 1
query = "Bad News"
print(ranksearch(query, limit))

[davidmezzetti]

No problem, glad I could help.

For reference, a laptop I have is 5+ years old, has a quad core CPU and a 8 GB GPU with 1920 CUDA cores. Modest specs compared to the most recent hardware. I’ve run benchmarks on this hardware below just to give you an idea of what you should expect.

import timeit
from txtai.embeddings import Embeddings
from txtai.pipeline import Similarity

# d is a list of 100 text elements (same count as code above when limit=10)

# GPU
similarity = Similarity("valhalla/distilbart-mnli-12-3")
timeit.timeit(lambda: similarity("query", d), number=25) / 25
# 1.55s per call

embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2"})
timeit.timeit(lambda: embeddings.similarity("query", d), number=25) / 25
# 0.45s per call

embeddings = Embeddings({"path": "sentence-transformers/paraphrase-MiniLM-L3-v2"})
timeit.timeit(lambda: embeddings.similarity("query", d), number=25) / 25
# 0.30s per call

# CPU only
similarity = Similarity("valhalla/distilbart-mnli-12-3")
timeit.timeit(lambda: similarity("query", d), number=25) / 25
# 19.16s per call

embeddings = Embeddings({"path": "sentence-transformers/all-MiniLM-L6-v2"})
timeit.timeit(lambda: embeddings.similarity("query", d), number=25) / 25
# 1.18s per call

embeddings = Embeddings({"path": "sentence-transformers/paraphrase-MiniLM-L3-v2"})
timeit.timeit(lambda: embeddings.similarity("query", d), number=25) / 25
# 0.66s per call

GPU prices have really come down lately. A RTX 3060 is ~$500 and there are RTX 3090s out there for around $1,100. A year ago those were 2.5-3x more expensive.

A RTX 3060 has 3,584 CUDA cores with 12GB of memory and a RTX 3090 has 10,496 with 24GB of memory. The elapsed time per call would be much lower on either of those. Server class NVIDIA GPUs are typically Quadro, V100, A100.

[Xyphius]

That does make quite the difference. Setting the limit to ten yields results around 4.3 seconds.

I will look into setting up something with respect to GPU processing. I greatly appreciate your help.