pinecone.data.features.inference.repl_overrides

  1from pinecone.core.openapi.inference.models import (
  2    EmbeddingsList as OpenAPIEmbeddingsList,
  3    SparseEmbedding as OpenAPISparseEmbedding,
  4    DenseEmbedding as OpenAPIDenseEmbedding,
  5    Document as OpenAPIDocument,
  6    RerankResult as OpenAPIRerankResult,
  7    RankedDocument as OpenAPIRankedDocument,
  8)
  9
 10
 11def present_list(mylist) -> str:
 12    if not mylist:
 13        return "[]"
 14    if len(mylist) <= 5:
 15        # Show list as usual when fewer than 5 items.
 16        # This number is arbitrary and can be adjusted
 17        # but it seems silly to show the abbreviated
 18        # message with (2  more) or whatever when the
 19        # number of items is so small and it's no problem
 20        # to display the real values.
 21        return f"[{', '.join(repr(x) for x in mylist)}]"
 22    first_part = ", ".join(repr(x) for x in mylist[:2])
 23    last_part = ", ".join(repr(x) for x in mylist[-2:])
 24    formatted_values = f"[{first_part}, ..., {last_part}]"
 25    return formatted_values
 26
 27
 28def sparse_embedding_to_str(self: OpenAPISparseEmbedding):
 29    formatted_sparse_values = present_list(self.sparse_values)
 30    formatted_sparse_indices = present_list(self.sparse_indices)
 31    formatted_sparse_tokens = present_list(self.sparse_tokens)
 32    parts = [
 33        "'vector_type': sparse",
 34        f"'sparse_values': {formatted_sparse_values}",
 35        f"'sparse_indices': {formatted_sparse_indices}",
 36    ]
 37    if self.sparse_tokens:
 38        parts.append(f"'sparse_tokens': {formatted_sparse_tokens}")
 39    return "{{{}}}".format(", ".join(parts))
 40
 41
 42def dense_embedding_to_str(self: OpenAPIDenseEmbedding):
 43    formatted_values = present_list(self.values)
 44    parts = ["'vector_type': dense", f"'values': {formatted_values}"]
 45    return "{{{}}}".format(", ".join(parts))
 46
 47
 48def embedding_list_to_string(self: OpenAPIEmbeddingsList):
 49    if len(self.data) == 0:
 50        formatted_embeddings = "[]"
 51    elif len(self.data) <= 5:
 52        formatted_embeddings = (
 53            "[\n    " + ",\n    ".join([embedding.to_str() for embedding in self.data]) + "\n  ]"
 54        )
 55    else:
 56        omitted_msg = f"... ({len(self.data) - 4} more embeddings) ..."
 57        items_to_show = (
 58            [e.to_str() for e in self.data[:2]]
 59            + [omitted_msg]
 60            + [e.to_str() for e in self.data[-2:]]
 61        )
 62        formatted_embeddings = "[\n    " + ",\n    ".join(items_to_show) + "\n  ]"
 63
 64    return """EmbeddingsList(
 65  model='{self.model}',
 66  vector_type='{self.vector_type}',
 67  data={formatted_embeddings},
 68  usage={self.usage}
 69)""".format(self=self, formatted_embeddings=formatted_embeddings)
 70
 71
 72def rank_result_document_to_str(self: OpenAPIDocument):
 73    return ",\n        ".join([f"{k}={repr(v)}" for k, v in self.to_dict().items()])
 74
 75
 76def rank_result_item_to_str(self: OpenAPIRankedDocument):
 77    if self.document is None:
 78        document_str = ""
 79    else:
 80        document_str = f"document={{\n        {self.document}\n    }}"
 81
 82    return f"{{\n    index={self.index},\n    score={self.score},\n    {document_str}\n  }}"
 83
 84
 85def rerank_result_to_str(self: OpenAPIRerankResult):
 86    if len(self.data) == 0:
 87        formatted_rerank = "[]"
 88    else:
 89        formatted_rerank = (
 90            "[" + ",".join([rank_result_item.to_str() for rank_result_item in self.data]) + "]"
 91        )
 92
 93    return f"""RerankResult(
 94  model='{self.model}',
 95  data={formatted_rerank},
 96  usage={self.usage}
 97)"""
 98
 99
100def install_repl_overrides():
101    OpenAPIDenseEmbedding.to_str = dense_embedding_to_str
102    OpenAPISparseEmbedding.to_str = sparse_embedding_to_str
103    OpenAPIEmbeddingsList.to_str = embedding_list_to_string
104    OpenAPIDocument.to_str = rank_result_document_to_str
105    OpenAPIRankedDocument.to_str = rank_result_item_to_str
106    OpenAPIRerankResult.to_str = rerank_result_to_str