diff --git a/README.md b/README.md index e0f427be..e019d981 100644 --- a/README.md +++ b/README.md @@ -487,7 +487,7 @@ UpdateResponse updateResponse = index.update("v1", values, "example-namespace"); # Collections -Collections fall under data plane operations. +Collections fall under control plane operations. ## Create collection @@ -545,8 +545,8 @@ Pinecone pinecone = new Pinecone.Builder("PINECONE_API_KEY").build(); pinecone.deleteCollection("example-collection"); ``` -## Inference - +# Inference +## Embed The Pinecone SDK now supports creating embeddings via the [Inference API](https://docs.pinecone.io/guides/inference/understanding-inference). ```java @@ -583,6 +583,66 @@ EmbeddingsList embeddings = inference.embed(embeddingModel, parameters, inputs); List embeddedData = embeddings.getData(); ``` +## Rerank +The following example shows how to rerank items according to their relevance to a query. + +```java +import io.pinecone.clients.Inference; +import io.pinecone.clients.Pinecone; +import org.openapitools.inference.client.model.RerankResult; + +import java.util.*; + +... + +// The model to use for reranking +String model = "bge-reranker-v2-m3"; + +// The query to rerank documents against +String query = "The tech company Apple is known for its innovative products like the iPhone."; + +// Add the documents to rerank +List> documents = new ArrayList<>(); +Map doc1 = new HashMap<>(); +doc1.put("id", "vec1"); +doc1.put("my_field", "Apple is a popular fruit known for its sweetness and crisp texture."); +documents.add(doc1); + +Map doc2 = new HashMap<>(); +doc2.put("id", "vec2"); +doc2.put("my_field", "Many people enjoy eating apples as a healthy snack."); +documents.add(doc2); + +Map doc3 = new HashMap<>(); +doc3.put("id", "vec3"); +doc3.put("my_field", "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."); +documents.add(doc3); + +Map doc4 = new HashMap<>(); +doc4.put("id", "vec4"); +doc4.put("my_field", "An apple a day keeps the doctor away, as the saying goes."); +documents.add(doc4); + +// The fields to rank the documents by. If not provided, the default is "text" +List rankFields = Arrays.asList("my_field"); + +// The number of results to return sorted by relevance. Defaults to the number of inputs +int topN = 2; + +// Whether to return the documents in the response +boolean returnDocuments = true; + +// Additional model-specific parameters for the reranker +Map parameters = new HashMap<>(); +parameters.put("truncate", "END"); + +// Send ranking request +RerankResult result = inference.rerank(model, query, documents, rankFields, topN, returnDocuments, parameters); + +// Get ranked data +System.out.println(result.getData()); +``` + ## Examples - The data and control plane operation examples can be found in `io/pinecone/integration` folder. \ No newline at end of file diff --git a/src/integration/java/io/pinecone/integration/inference/RerankTest.java b/src/integration/java/io/pinecone/integration/inference/RerankTest.java new file mode 100644 index 00000000..6ef266b3 --- /dev/null +++ b/src/integration/java/io/pinecone/integration/inference/RerankTest.java @@ -0,0 +1,95 @@ +package io.pinecone.integration.inference; + +import io.pinecone.clients.Inference; +import io.pinecone.clients.Pinecone; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.openapitools.inference.client.ApiException; +import org.openapitools.inference.client.model.RerankResult; + +import java.util.*; + +import static org.junit.jupiter.api.Assertions.assertNotNull; + +public class RerankTest { + private static final Pinecone pinecone = new Pinecone + .Builder(System.getenv("PINECONE_API_KEY")) + .withSourceTag("pinecone_test") + .build(); + private static final Inference inference = pinecone.getInferenceClient(); + + @Test + public void testRerank() throws ApiException { + String model = "bge-reranker-v2-m3"; + String query = "The tech company Apple is known for its innovative products like the iPhone."; + List> documents = new ArrayList<>(); + + Map doc1 = new HashMap<>(); + doc1.put("id", "vec1"); + doc1.put("my_field", "Apple is a popular fruit known for its sweetness and crisp texture."); + documents.add(doc1); + + Map doc2 = new HashMap<>(); + doc2.put("id", "vec2"); + doc2.put("my_field", "Many people enjoy eating apples as a healthy snack."); + documents.add(doc2); + + Map doc3 = new HashMap<>(); + doc3.put("id", "vec3"); + doc3.put("my_field", "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."); + documents.add(doc3); + + Map doc4 = new HashMap<>(); + doc4.put("id", "vec4"); + doc4.put("my_field", "An apple a day keeps the doctor away, as the saying goes."); + documents.add(doc4); + + List rankFields = Arrays.asList("my_field"); + int topN = 2; + boolean returnDocuments = true; + Map parameters = new HashMap<>(); + parameters.put("truncate", "END"); + + RerankResult result = inference.rerank(model, query, documents, rankFields, topN, returnDocuments, parameters); + + assertNotNull(result); + Assertions.assertEquals(result.getData().size(), topN); + Assertions.assertEquals(result.getData().get(0).getIndex(), 2); + Assertions.assertEquals(result.getData().get(0).getDocument().get("my_field"), doc3.get("my_field")); + Assertions.assertEquals(result.getData().size(), 2); + } + + @Test + public void testRerankWithRequiredParameters() throws ApiException { + String model = "bge-reranker-v2-m3"; + String query = "The tech company Apple is known for its innovative products like the iPhone."; + List> documents = new ArrayList<>(); + + Map doc1 = new HashMap<>(); + doc1.put("id", "vec1"); + doc1.put("text", "Apple is a popular fruit known for its sweetness and crisp texture."); + documents.add(doc1); + + Map doc2 = new HashMap<>(); + doc2.put("id", "vec2"); + doc2.put("text", "Many people enjoy eating apples as a healthy snack."); + documents.add(doc2); + + Map doc3 = new HashMap<>(); + doc3.put("id", "vec3"); + doc3.put("text", "Apple Inc. has revolutionized the tech industry with its sleek designs and user-friendly interfaces."); + documents.add(doc3); + + Map doc4 = new HashMap<>(); + doc4.put("id", "vec4"); + doc4.put("text", "An apple a day keeps the doctor away, as the saying goes."); + documents.add(doc4); + + RerankResult result = inference.rerank(model, query, documents); + + assertNotNull(result); + Assertions.assertEquals(result.getData().size(), documents.size()); + Assertions.assertEquals(result.getData().get(0).getIndex(), 2); + Assertions.assertEquals(result.getData().get(0).getDocument().get("text"), doc3.get("text")); + } +} diff --git a/src/main/java/io/pinecone/clients/Inference.java b/src/main/java/io/pinecone/clients/Inference.java index b8d5761f..703db2c1 100644 --- a/src/main/java/io/pinecone/clients/Inference.java +++ b/src/main/java/io/pinecone/clients/Inference.java @@ -60,6 +60,65 @@ public EmbeddingsList embed(String model, Map parameters, List> documents) throws ApiException { + return rerank(model, + query, + documents, + Arrays.asList("text"), + documents.size(), + true, + new HashMap<>()); + } + + /** + * Reranks a list of documents based on the relevance to a query using the specified model with additional options. + * + * @param model The model to be used for reranking the documents. + * @param query The query string to rank the documents against. + * @param documents A list of maps representing the documents to be ranked. + * Each map should contain document attributes, such as "text". + * @param rankFields A list of fields in the documents to be used for ranking, typically "text". + * @param topN The number of top-ranked documents to return. + * @param returnDocuments Whether to return the documents along with the ranking scores. + * @param parameters A map containing additional model-specific parameters for reranking. + * @return RerankResult containing the ranked documents and their scores. + * @throws ApiException If the API call fails, an ApiException is thrown. + */ + public RerankResult rerank(String model, + String query, + List> documents, + List rankFields, + int topN, + boolean returnDocuments, + Map parameters) throws ApiException { + RerankRequest rerankRequest = new RerankRequest(); + + rerankRequest + .model(model) + .query(query) + .documents(documents) + .rankFields(rankFields) + .topN(topN) + .returnDocuments(returnDocuments) + .putAdditionalProperty("parameters", parameters); + + return inferenceApi.rerank(rerankRequest); + } + /** * Converts a list of input strings to EmbedRequestInputsInner objects. *