​​Reinventing Multi-modal Search With MongoDB and Anyscale

1
{
2
  "fields": [
3
    {
4
      "numDimensions": 1024,
5
      "similarity": "cosine",
6
      "type": "vector",
7
      "path": "description_embedding"
8
    },
9
    {
10
      "numDimensions": 1024,
11
      "similarity": "cosine",
12
      "type": "vector",
13
      "path": "name_embedding"
14
    },
15
    {
16
      "type": "filter",
17
      "path": "category"
18
    },
19
    {
20
      "type": "filter",
21
      "path": "season"
22
    },
23
    {
24
      "type": "filter",
25
      "path": "color"
26
    },
27
    {
28
      "type": "filter",
29
      "path": "rating"
30
    },
31
    {
32
      "type": "filter",
33
      "path": "price"
34
    }
35
  ]
36
}

1
{
2
  "mappings": {
3
    "dynamic": false,
4
    "fields": {
5
      "name": {
6
        "type": "string",
7
        "analyzer": "lucene.standard"
8
      }
9
    }
10
  }
11
}

1
ds = ray.data.read_csv(path, ...)
2
ds = (
3
ds.map_batches(download_images, concurrency=num_download_image_workers,
4
num_cpus=4)
5
.filter(lambda x: bool(x["img"]))
6
.map(LargestCenterSquare(size=336))
7
.map(gen_description_prompt)
8
.materialize()
9
)

1
def gen_description_prompt(row: dict[str, Any]) -> dict[str, Any]:
2
  title = row["name"]
3
  row["description_prompt"] = "<image>" * 1176 + (
4
    f"\nUSER: Generate an ecommerce product description given the image
5
    and this title: {title}."
6
    "Make sure to include information about the color of the product in
7
    the description."
8
    "\nASSISTANT:"
9
  )
10
  return row

1
max_input_tokens = (
2
ds.map_batches(
3
LlaVAMistralTokenizer,
4
fn_kwargs={
5
"input": "description_prompt",
6
"output": "description_prompt_tokens",
7
},
8
concurrency=num_llava_tokenizer_workers,
9
num_cpus=1,
10
)
11
.select_columns(["description_prompt_tokens"])
12
.map(compute_num_tokens, fn_kwargs={"col": "description_prompt_tokens"})
13
.max(on="num_tokens")
14
)

1
max_output_tokens = 256
2
max_model_length = max_input_tokens + max_output_tokens

1
ds = ds.map_batches(
2
LlaVAMistral,
3
fn_constructor_kwargs={
4
"max_model_len": max_model_length,
5
"max_tokens": max_output_tokens,
6
"max_num_seqs": 400,
7
},
8
fn_kwargs={"col": "description_prompt"},
9
batch_size=llava_model_batch_size,
10
num_gpus=1,
11
concurrency=num_llava_model_workers,
12
accelerator_type=llava_model_accelerator_type,
13
)

1
class LlaVAMistral:
2
  def __init__(self, max_model_len: int, ...):
3
    self.llm = LLM(...)
4
5
  def __call__(self, batch: dict[str, np.ndarray], col: str) -> dict[str, np.ndarray]:
6
    prompts = batch[col]
7
    images = batch["img"]
8
    responses = self.llm.generate(
9
      [
10
        {
11
          "prompt": prompt,
12
          "multi_modal_data": ImagePixelData(image),
13
        }
14
        for prompt, image in zip(prompts, images)
15
      ],
16
      sampling_params=self.sampling_params,
17
    )
18
    batch["description"] = [resp.outputs[0].text for resp in responses] # type: ignore
19
    return batch

1
classifiers: dict[str, Any] = {
2
  "category": {
3
    "classes": ["Tops", "Bottoms", "Dresses", "Footwear", "Accessories"],
4
    "prompt_template": (
5
      "Given the title of this product: {title} and "
6
      "the description: {description}, what category does it belong to?"
7
      "Chose from the following categories: {classes_str}. "
8
      "Return the category that best fits the product. Only return the"
9
      "category name and nothing else."
10
    ),
11
    "prompt_constructor": construct_prompt_classifier,
12
  },
13
  "season": {
14
    "classes": ["Summer", "Winter", "Spring", "Fall"],
15
    "prompt_template": ...,
16
    "prompt_constructor": construct_prompt_classifier,
17
  },
18
  "color": {
19
    ...
20
  }
21
}

1
for classifier, classifier_spec in classifiers.items():
2
  ds = (
3
    ds.map(
4
      classifier_spec["prompt_constructor"],
5
      fn_kwargs={
6
        "prompt_template": classifier_spec["prompt_template"],
7
        "classes": classifier_spec["classes"],
8
        "col": classifier,
9
      },
10
    )
11
    .map_batches(
12
      MistralTokenizer,
13
      fn_kwargs={
14
        "input": f"{classifier}_prompt",
15
        "output": f"{classifier}_prompt_tokens",
16
      },
17
      concurrency=num_mistral_tokenizer_workers_per_classifier,
18
      num_cpus=1,
19
    )
20
    .materialize()
21
  )

1
for classifier, classifier_spec in classifiers.items():
2
  max_input_tokens = (
3
    ds.select_columns([f"{classifier}_prompt_tokens"])
4
    .map(compute_num_tokens, fn_kwargs={"col": f"{classifier}_prompt_tokens"})
5
    .max(on="num_tokens")
6
  )
7
  max_output_tokens = classifier_spec["max_output_tokens"]
8
  max_model_length = max_input_tokens + max_output_tokens
9
  classifier_spec["max_model_length"] = max_model_length

1
for classifier, classifier_spec in classifiers.items():
2
  ds = (
3
    ds.map_batches(
4
      MistralvLLM,
5
      ...,
6
      batch_size=80,
7
      num_gpus=num_mistral_workers_per_classifier,
8
      concurrency=1,
9
      accelerator_type=NVIDIA_TESLA_A10G,
10
    )
11
    .map(
12
      MistralDeTokenizer,
13
      fn_kwargs={"key": f"{classifier}_response"},
14
      concurrency=num_detokenizers_per_classifier,
15
      num_cpus=1,
16
    )
17
    .map(
18
      clean_response,
19
      fn_kwargs={
20
        "classes": classifier_spec["classes"],
21
        "response_col": f"{classifier}_response",
22
      },
23
    )
24
  )

1
ds = ds.map_batches(
2
  EmbedderSentenceTransformer,
3
  fn_kwargs={"cols": ["name", "description"]},
4
  batch_size=80,
5
  num_gpus=1,
6
  concurrency=num_embedder_workers,
7
  accelerator_type=NVIDIA_TESLA_A10G,
8
)

1
class EmbedderSentenceTransformer:
2
  def __init__(self, model: str = "thenlper/gte-large"):
3
    self.model = SentenceTransformer(model, device="cuda")
4
5
  def __call__(self, batch: dict[str, np.ndarray], cols: list[str]) -> dict[str, np.ndarray]:
6
    for col in cols:
7
      batch[f"{col}_embedding"] = self.model.encode( # type: ignore
8
        batch[col].tolist(), batch_size=len(batch[col])
9
      )
10
    return batch

1
mongo_bulk_op: Type[MongoBulkInsert] | Type[MongoBulkUpdate]
2
if mode == "first_run":
3
  mongo_bulk_op = MongoBulkInsert
4
elif mode == "update":
5
  mongo_bulk_op = MongoBulkUpdate
6
7
(
8
  ds.map_batches(update_record)
9
  .map_batches(
10
    mongo_bulk_op,
11
    fn_constructor_kwargs={
12
      "db": db_name,
13
      "collection": collection_name,
14
    },
15
    batch_size=db_update_batch_size,
16
    concurrency=num_db_workers,
17
    num_cpus=0.1,
18
    batch_format="pandas",
19
  )
20
  .materialize()
21
)

1
class MongoBulkUpdate:
2
  def __init__(self, db: str, collection: str) -> None:
3
    client = MongoClient(os.environ["DB_CONNECTION_STRING"])
4
    self.collection = client[db][collection]
5
6
  def __call__(self, batch_df: pd.DataFrame) -> dict[str, np.ndarray]:
7
    docs = batch_df.to_dict(orient="records")
8
    bulk_ops = [
9
      UpdateOne(filter={"_id": doc["_id"]}, update={"$set": doc}, upsert=True)
10
      for doc in docs
11
    ]
12
    self.collection.bulk_write(bulk_ops)
13
    return {}

1
name: enrich-data-and-upsert
2
entrypoint: python cli.py ...
3
working_dir: .
4
requirements: requirements.txt
5
env_vars:
6
  DB_CONNECTION_STRING: <your mongodb connection string>
7
  HF_TOKEN: <your huggingface token>
8
compute_config:
9
  cloud: "Anyscale Cloud"
10
  head_node:
11
    instance_type: m5.8xlarge
12
  worker_nodes:
13
    - instance_type: m5.8xlarge
14
      min_nodes: 0
15
      max_nodes: 10
16
    - instance_type: g5.xlarge
17
      min_nodes: 0
18
      max_nodes: 40

1
anyscale job submit -f job.yaml

1
fastapi = FastAPI()
2
3
@deployment
4
@ingress(fastapi)
5
class QueryApplication:
6
  def __init__(self, query_legacy: QueryLegacySearch, query_ai_enabled: QueryAIEnabledSearch):
7
    self.query_legacy = query_legacy
8
    self.query_ai_enabled = query_ai_enabled
9
10
  @fastapi.get("/legacy")
11
  async def query_legacy_search(
12
    self,
13
    text_search: str,
14
    min_price: int,
15
    max_price: int,
16
    min_rating: float,
17
    num_results: int,
18
  ):
19
    return await self.query_legacy.run.remote(...)
20
21
  @fastapi.get("/ai_enabled")
22
  async def query_ai_enabled_search(
23
    self,
24
    text_search: str,
25
    min_price: int,
26
    max_price: int,
27
    min_rating: float,
28
    categories: list[str],
29
    colors: list[str],
30
    seasons: list[str],
31
    num_results: int,
32
    embedding_column: str,
33
    search_type: list[str],
34
  ):
35
    logger = logging.getLogger("ray.serve")
36
    logger.setLevel(logging.DEBUG)
37
    logger.debug(f"Running query_ai_enabled_search with {locals()=}")
38
    return await self.query_ai_enabled.run.remote(...)

1
@deployment
2
class QueryAIEnabledSearch:
3
  def __init__(
4
    self,
5
    embedding_model: DeploymentHandle,
6
    database_name: str,
7
    collection_name: str,
8
  ) -> None:
9
    self.client = Async
10
11
IOMotorClient(os.environ["DB_CONNECTION_STRING"])
12
    self.embedding_model = embedding_model
13
    self.database_name = database_name
14
    self.collection_name = collection_name
15
16
  async def run(
17
    self,
18
    text_search: str,
19
    min_price: int,
20
    max_price: int,
21
    min_rating: float,
22
    categories: list[str],
23
    colors: list[str],
24
    seasons: list[str],
25
    n: int,
26
    search_type: set[str],
27
    vector_search_index_name: str = "vector_search_index",
28
    vector_search_path: str = "description_embedding",
29
    text_search_index_name: str = "lexical_text_search_index",
30
    vector_penalty: int = 1,
31
    full_text_penalty: int = 10,
32
  ):
33
    db = self.client[self.database_name]
34
    collection = db[self.collection_name]
35
    pipeline = []
36
    if text_search.strip():
37
      if "vector" in search_type:
38
        embedding = await self.embedding_model.compute_embedding.remote(text_search)
39
      is_hybrid = search_type == {"vector", "lexical"}
40
      if is_hybrid:
41
        pipeline.extend(hybrid_search(...))
42
      elif search_type == {"vector"}:
43
        pipeline.extend(vector_search(...))
44
      elif search_type == {"lexical"}:
45
        pipeline.extend(lexical_search(...))
46
      pipeline.extend(match_on_metadata(...))
47
    else:
48
      pipeline = match_on_metadata(...)
49
50
    records = collection.aggregate(pipeline)
51
    records = [record async for record in records]
52
    results = [(record["img"], record["name"]) for record in records]
53
    return results

1
def vector_search(
2
  vector_search_index_name: str,
3
  vector_search_path: str,
4
  embedding: list[float],
5
  n: int,
6
  min_price: int,
7
  max_price: int,
8
  min_rating: float,
9
  categories: list[str],
10
  colors: list[str],
11
  seasons: list[str],
12
  cosine_score_threshold: float = 0.92,
13
) -> list[dict]:
14
  return [
15
    {
16
      "$vectorSearch": {
17
        "index": vector_search_index_name,
18
        "path": vector_search_path,
19
        "queryVector": embedding.tolist(),
20
        "numCandidates": 100,
21
        "limit": n,
22
        "filter": {
23
          "price": {"$gte": min_price, "$lte": max_price},
24
          "rating": {"$gte": min_rating},
25
          "category": {"$in": categories},
26
          "color": {"$in": colors},
27
          "season": {"$in": seasons},
28
        },
29
      },
30
    },
31
    {
32
      "$project": {
33
        "img": 1,
34
        "name": 1,
35
        "score": {"$meta": "vectorSearchScore"},
36
      }
37
    },
38
    {"$match": {"score": {"$gte": cosine_score_threshold}}},
39
  ]

1
def lexical_search(text_search: str) -> list[dict]:
2
  return [
3
    {
4
      "$search": {
5
        "index": "lexical_text_search_index",
6
        "text": {
7
          "query": text_search,
8
          "path": "name",
9
        },
10
      }
11
    }
12
  ]

1
def hybrid_search(
2
  collection_name: str,
3
  ...
4
) -> list[dict]:
5
  # 1. Perform vector search
6
  vector_search_stages = vector_search(...)
7
  convert_vector_rank_to_score_stages = convert_rank_to_score(
8
    score_name="vs_score", score_penalty=vector_penalty
9
  )
10
11
  # 2. Perform lexical search
12
  lexical_search_stages = lexical_search(text_search=text_search, text_search_index_name=text_search_index_name)
13
  post_filter_stages = match_on_metadata(...)
14
  convert_text_rank_to_score_stages = convert_rank_to_score(
15
    score_name="fts_score", score_penalty=full_text_penalty
16
  )
17
18
  # 3. Rerank by combined score
19
  rerank_stages = rerank_by_combined_score(
20
    vs_score_name="vs_score", fts_score_name="fts_score", n=n
21
  )
22
23
  # 4. Put it all together
24
  return [
25
    *vector_search_stages,
26
    *convert_vector_rank_to_score_stages,
27
    {
28
      "$unionWith": {
29
        "coll": collection_name,
30
        "pipeline": [
31
          *lexical_search_stages,
32
          *post_filter_stages,
33
          *convert_text_rank_to_score_stages,
34
        ],
35
      }
36
    },
37
    *rerank_stages,
38
  ]

1
@deployment
2
class EmbeddingModel:
3
  def __init__(self, model: str = "thenlper/gte-large") -> None:
4
    self.model = SentenceTransformer(model)
5
6
  async def compute_embedding(self, text: str) -> list[float]:
7
    loop = asyncio.get_event_loop()
8
    return await loop.run_in_executor(None, lambda: self.model.encode(text))

1
name: mongo-multi-modal-search-v2
2
applications:
3
  - name: frontend
4
    route_prefix: /
5
    import_path: frontend:app
6
  - name: backend
7
    route_prefix: /backend
8
    import_path: backend:app
9
query_auth_token_enabled: false
10
requirements: requirements.txt

1
anyscale service deploy -f app.yaml