TheoKanning / openai-java

OpenAI Api Client in Java
MIT License
4.75k stars 1.19k forks source link

Feature: Add dimension paramater in embedding request #470

Closed prabhupant closed 7 months ago

prabhupant commented 7 months ago

Thanks for submitting a pull request! Please check CONTRIBUTING.md for style guidelines.

Changes

Currently the EmbeddingRequest does not have a parameter dimensions. So the default value that OpenAPI uses, that is 1536, is set for all the embedding request. If you want to generate embeddings with a reduced dimensions, currently this is not possible.

With the recent launch of new v3 embedding model, text-small-3-embedding performs better than ada-002 even with reduced dimension. So this feature can be essential for someone who is looking to save on their storage cost by using a reduced dimension without compromising on the quality of the generated vector embedding

New API Checklist

See CONTRIBUTING.md for more info.

  1. [x] Documentation for every variable
  2. [ ] Class-level documentation
  3. [ ] POJO JSON parsing tests
  4. [ ] Service integration tests
VladBrok commented 5 months ago

@prabhupant hi. Why did you close it :(

prabhupant commented 5 months ago

@VladBrok I inherited the base class and overrode the function to create my own method. It was something like this

This was the interface that extended the one provided in this package

import com.theokanning.openai.client.OpenAiApi;
import com.theokanning.openai.embedding.EmbeddingResult;
import io.reactivex.Single;
import retrofit2.http.Body;
import retrofit2.http.POST;

/** Interface for using embedding model with dimension size. */
public interface OpenAiApiExt extends OpenAiApi {

  @POST("/v1/embeddings")
  Single<EmbeddingResult> createEmbeddingsV2(@Body EmbeddingV2Request var1);
}

And then here is the implementation

import com.fasterxml.jackson.databind.ObjectMapper;
import com.theokanning.openai.embedding.EmbeddingResult;
import com.theokanning.openai.service.OpenAiService;
import java.time.Duration;
import okhttp3.OkHttpClient;
import retrofit2.Retrofit;

/** Wrapper class of OpenAiService that enables request with an updated request body. */
public class OpenAiServiceV2 extends OpenAiService {

  private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(10);

  private final OpenAiApiExt apiV2;

  /**
   * Creates a new OpenAiService that wraps OpenAiApi.
   *
   * @param token OpenAi token
   */
  public OpenAiServiceV2(final String token) {
    super(token);

    ObjectMapper mapper = defaultObjectMapper();
    OkHttpClient client = defaultClient(token, DEFAULT_TIMEOUT);
    Retrofit retrofit = defaultRetrofit(client, mapper);

    this.apiV2 = retrofit.create(OpenAiApiExt.class);
  }

  public EmbeddingResult createEmbeddingsV2(EmbeddingV2Request request) {
    return execute(apiV2.createEmbeddingsV2(request));
  }
}

And here is the EmbeddingV2Request model

import java.util.List;
import lombok.AllArgsConstructor;
import lombok.Builder;
import lombok.Data;
import lombok.NoArgsConstructor;
import lombok.NonNull;

@Builder
@NoArgsConstructor
@AllArgsConstructor
@Data
public class EmbeddingV2Request {

  String model;

  @NonNull List<String> input;

  String user;

  Integer dimensions;
}

Hope you find it helpful :)

VladBrok commented 5 months ago

@prabhupant thank you! looks cool :muscle: