code: def semantic_search(
      query_embeddings: Tensor,
      corpus_embeddings: Tensor,
      query_chunk_size: int = 100,
      corpus_chunk_size: int = 500000,
      top_k: int = 10,
      score_function: Callable[[Tensor, Tensor], Tensor] = cos_sim,
  ) -> List[List[Dict[str, Union[int, float]]]]:
      """
      This function performs a cosine similarity search between a list of query embeddings  and a list of corpus embeddings.
      It can be used for Information Retrieval / Semantic Search for corpora up to about 1 Million entries.
  
      :param query_embeddings: A 2 dimensional tensor with the query embeddings.
      :param corpus_embeddings: A 2 dimensional tensor with the corpus embeddings.
      :param query_chunk_size: Process 100 queries simultaneously. Increasing that value increases the speed, but requires more memory.
      :param corpus_chunk_size: Scans the corpus 100k entries at a time. Increasing that value increases the speed, but requires more memory.
      :param top_k: Retrieve top k matching entries.
      :param score_function: Function for computing scores. By default, cosine similarity.
      :return: Returns a list with one entry for each query. Each entry is a list of dictionaries with the keys 'corpus_id' and 'score', sorted by decreasing cosine similarity scores.
      """
  
      if isinstance(query_embeddings, (np.ndarray, np.generic)):
          query_embeddings = torch.from_numpy(query_embeddings)
      elif isinstance(query_embeddings, list):
          query_embeddings = torch.stack(query_embeddings)
  
      if len(query_embeddings.shape) == 1:
          query_embeddings = query_embeddings.unsqueeze(0)
  
      if isinstance(corpus_embeddings, (np.ndarray, np.generic)):
          corpus_embeddings = torch.from_numpy(corpus_embeddings)
      elif isinstance(corpus_embeddings, list):
          corpus_embeddings = torch.stack(corpus_embeddings)
  
      # Check that corpus and queries are on the same device
      if corpus_embeddings.device != query_embeddings.device:
          query_embeddings = query_embeddings.to(
              corpus_embeddings.device
          )
  
      queries_result_list = [[] for _ in range(len(query_embeddings))]
  
      for query_start_idx in range(
          0, len(query_embeddings), query_chunk_size
      ):
          # Iterate over chunks of the corpus
          for corpus_start_idx in range(
              0, len(corpus_embeddings), corpus_chunk_size
          ):
              # Compute cosine similarities
              cos_scores = score_function(
                  query_embeddings[
                      query_start_idx : query_start_idx
                      + query_chunk_size
                  ],
                  corpus_embeddings[
                      corpus_start_idx : corpus_start_idx
                      + corpus_chunk_size
                  ],
              )
  
              # Get top-k scores
              (
                  cos_scores_top_k_values,
                  cos_scores_top_k_idx,
              ) = torch.topk(
                  cos_scores,
                  min(top_k, len(cos_scores[0])),
                  dim=1,
                  largest=True,
                  sorted=False,
              )
              cos_scores_top_k_values = (
                  cos_scores_top_k_values.cpu().tolist()
              )
              cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist()
  
              for query_itr in range(len(cos_scores)):
                  for sub_corpus_id, score in zip(
                      cos_scores_top_k_idx[query_itr],
                      cos_scores_top_k_values[query_itr],
                  ):
                      corpus_id = corpus_start_idx + sub_corpus_id
                      query_id = query_start_idx + query_itr
                      if len(queries_result_list[query_id]) < top_k:
                          heapq.heappush(
                              queries_result_list[query_id],
                              (score, corpus_id),
                          )  # heaqp tracks the quantity of the first element in the tuple
                      else:
                          heapq.heappushpop(
                              queries_result_list[query_id],
                              (score, corpus_id),
                          )
  
      # change the data format and sort
      for query_id in range(len(queries_result_list)):
          for doc_itr in range(len(queries_result_list[query_id])):
              score, corpus_id = queries_result_list[query_id][doc_itr]
              queries_result_list[query_id][doc_itr] = {
                  "corpus_id": corpus_id,
                  "score": score,
              }
          queries_result_list[query_id] = sorted(
              queries_result_list[query_id],
              key=lambda x: x["score"],
              reverse=True,
          )
  
      return queries_result_list
  
imports: import functools
  import requests
  from torch import Tensor, device
  from typing import List, Callable
  from tqdm.autonotebook import tqdm
  import sys
  import importlib
  import os
  import torch
  import numpy
  import queue
  import logging
  from typing import Dict, Optional, Union
  from huggingface_hub import snapshot_download, hf_hub_download
  import heapq
label_desc: Use image node_id 10 for calling a global function with example usage: semantic_search(query_embeddings, corpus_embeddings, query_chunk_size, corpus_chunk_size, top_k, score_function) and returns: queries_result_list