import torch


def embed_queries_glide(queries, glide_model, glide_options, device):
  """Embed textual queries with the GLIDE text encoder."""

  tokens = []
  mask = []
  for query_mod in queries:
    # Create the caption tokens.
    view_tokens = glide_model.tokenizer.encode(query_mod)
    view_tokens, view_mask = glide_model.tokenizer.padded_tokens_and_mask(
        view_tokens, glide_options["text_ctx"])
    tokens.append(view_tokens)
    mask.append(view_mask)

  # Create the classifier-free guidance tokens (empty).
  uncond_tokens, uncond_mask = glide_model.tokenizer.padded_tokens_and_mask(
      [], glide_options["text_ctx"])
  uncond_tokens = [uncond_tokens] * len(queries)
  uncond_mask = [uncond_mask] * len(queries)

  # Pack the tokens together into model kwargs.
  model_kwargs = dict(
      tokens=torch.tensor(tokens + uncond_tokens, device=device),
      mask=torch.tensor(mask + uncond_mask, dtype=torch.bool, device=device),
  )
  return model_kwargs