Refactor PostgresML embedding options builder
- Deprecate the builder methods with the prefix `with` - Update the docs and references
This commit is contained in:
committed by
Mark Pollack
parent
95c32f05a3
commit
18eb4c3c61
@@ -33,6 +33,7 @@ import org.springframework.ai.postgresml.PostgresMlEmbeddingModel.VectorType;
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Thomas Vitale
|
||||
* @author Ilayaperumal Gopinathan
|
||||
*/
|
||||
@JsonInclude(Include.NON_NULL)
|
||||
public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
|
||||
@@ -116,26 +117,71 @@ public class PostgresMlEmbeddingOptions implements EmbeddingOptions {
|
||||
this.options = new PostgresMlEmbeddingOptions();
|
||||
}
|
||||
|
||||
public Builder transformer(String transformer) {
|
||||
this.options.setTransformer(transformer);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder vectorType(VectorType vectorType) {
|
||||
this.options.setVectorType(vectorType);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder kwargs(String kwargs) {
|
||||
this.options.setKwargs(ModelOptionsUtils.objectToMap(kwargs));
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder kwargs(Map<String, Object> kwargs) {
|
||||
this.options.setKwargs(kwargs);
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder metadataMode(MetadataMode metadataMode) {
|
||||
this.options.setMetadataMode(metadataMode);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #transformer(String)} )} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withTransformer(String transformer) {
|
||||
this.options.setTransformer(transformer);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #vectorType(VectorType)} )} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withVectorType(VectorType vectorType) {
|
||||
this.options.setVectorType(vectorType);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #kwargs(String)} )} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withKwargs(String kwargs) {
|
||||
this.options.setKwargs(ModelOptionsUtils.objectToMap(kwargs));
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #kwargs(Map)} )} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withKwargs(Map<String, Object> kwargs) {
|
||||
this.options.setKwargs(kwargs);
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated use {@link #metadataMode(MetadataMode)} )} instead.
|
||||
*/
|
||||
@Deprecated(forRemoval = true, since = "1.0.0-M5")
|
||||
public Builder withMetadataMode(MetadataMode metadataMode) {
|
||||
this.options.setMetadataMode(metadataMode);
|
||||
return this;
|
||||
|
||||
@@ -90,8 +90,8 @@ class PostgresMlEmbeddingModelIT {
|
||||
void embedWithPgVector() {
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("distilbert-base-uncased")
|
||||
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.transformer("distilbert-base-uncased")
|
||||
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.build(),
|
||||
true);
|
||||
embeddingModel.afterPropertiesSet();
|
||||
@@ -104,7 +104,7 @@ class PostgresMlEmbeddingModelIT {
|
||||
@Test
|
||||
void embedWithDifferentModel() {
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder().withTransformer("intfloat/e5-small").build(), true);
|
||||
PostgresMlEmbeddingOptions.builder().transformer("intfloat/e5-small").build(), true);
|
||||
embeddingModel.afterPropertiesSet();
|
||||
|
||||
float[] embed = embeddingModel.embed(new Document("Hello World!"));
|
||||
@@ -116,10 +116,10 @@ class PostgresMlEmbeddingModelIT {
|
||||
void embedWithKwargs() {
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("distilbert-base-uncased")
|
||||
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
|
||||
.withKwargs(Map.of("device", "cpu"))
|
||||
.withMetadataMode(MetadataMode.EMBED)
|
||||
.transformer("distilbert-base-uncased")
|
||||
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
|
||||
.kwargs(Map.of("device", "cpu"))
|
||||
.metadataMode(MetadataMode.EMBED)
|
||||
.build(),
|
||||
true);
|
||||
embeddingModel.afterPropertiesSet();
|
||||
@@ -134,8 +134,8 @@ class PostgresMlEmbeddingModelIT {
|
||||
void embedForResponse(String vectorType) {
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("distilbert-base-uncased")
|
||||
.withVectorType(VectorType.valueOf(vectorType))
|
||||
.transformer("distilbert-base-uncased")
|
||||
.vectorType(VectorType.valueOf(vectorType))
|
||||
.build(),
|
||||
true);
|
||||
embeddingModel.afterPropertiesSet();
|
||||
@@ -173,8 +173,8 @@ class PostgresMlEmbeddingModelIT {
|
||||
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("distilbert-base-uncased")
|
||||
.withVectorType(VectorType.PG_VECTOR)
|
||||
.transformer("distilbert-base-uncased")
|
||||
.vectorType(VectorType.PG_VECTOR)
|
||||
.build(),
|
||||
true);
|
||||
embeddingModel.afterPropertiesSet();
|
||||
@@ -211,10 +211,10 @@ class PostgresMlEmbeddingModelIT {
|
||||
// Override the default options in the request
|
||||
var request2 = new EmbeddingRequest(List.of("Hello World!", "Spring AI!", "LLM!"),
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("intfloat/e5-small")
|
||||
.withVectorType(VectorType.PG_ARRAY)
|
||||
.withMetadataMode(MetadataMode.EMBED)
|
||||
.withKwargs(Map.of("device", "cpu"))
|
||||
.transformer("intfloat/e5-small")
|
||||
.vectorType(VectorType.PG_ARRAY)
|
||||
.metadataMode(MetadataMode.EMBED)
|
||||
.kwargs(Map.of("device", "cpu"))
|
||||
.build());
|
||||
|
||||
embeddingResponse = embeddingModel.call(request2);
|
||||
|
||||
@@ -44,10 +44,10 @@ public class PostgresMlEmbeddingOptionsTests {
|
||||
@Test
|
||||
public void newOptions() {
|
||||
PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("intfloat/e5-small")
|
||||
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
|
||||
.withKwargs(Map.of("device", "cpu"))
|
||||
.transformer("intfloat/e5-small")
|
||||
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.metadataMode(org.springframework.ai.document.MetadataMode.ALL)
|
||||
.kwargs(Map.of("device", "cpu"))
|
||||
.build();
|
||||
|
||||
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
|
||||
@@ -72,8 +72,8 @@ public class PostgresMlEmbeddingOptionsTests {
|
||||
|
||||
// Partial override
|
||||
options = embeddingModel.mergeOptions(PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("intfloat/e5-small")
|
||||
.withKwargs(Map.of("device", "cpu"))
|
||||
.transformer("intfloat/e5-small")
|
||||
.kwargs(Map.of("device", "cpu"))
|
||||
.build());
|
||||
|
||||
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
|
||||
@@ -83,10 +83,10 @@ public class PostgresMlEmbeddingOptionsTests {
|
||||
|
||||
// Complete override
|
||||
options = embeddingModel.mergeOptions(PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("intfloat/e5-small")
|
||||
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.withMetadataMode(org.springframework.ai.document.MetadataMode.ALL)
|
||||
.withKwargs(Map.of("device", "cpu"))
|
||||
.transformer("intfloat/e5-small")
|
||||
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_VECTOR)
|
||||
.metadataMode(org.springframework.ai.document.MetadataMode.ALL)
|
||||
.kwargs(Map.of("device", "cpu"))
|
||||
.build());
|
||||
|
||||
assertThat(options.getTransformer()).isEqualTo("intfloat/e5-small");
|
||||
|
||||
@@ -77,9 +77,9 @@ For example to override the default model name for a specific request:
|
||||
EmbeddingResponse embeddingResponse = embeddingModel.call(
|
||||
new EmbeddingRequest(List.of("Hello World", "World is big and salvation is near"),
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("intfloat/e5-small")
|
||||
.withVectorType(VectorType.PG_ARRAY)
|
||||
.withKwargs(Map.of("device", "gpu"))
|
||||
.transformer("intfloat/e5-small")
|
||||
.vectorType(VectorType.PG_ARRAY)
|
||||
.kwargs(Map.of("device", "gpu"))
|
||||
.build()));
|
||||
----
|
||||
|
||||
@@ -148,10 +148,10 @@ var jdbcTemplate = new JdbcTemplate(dataSource); // your posgresml data source
|
||||
|
||||
PostgresMlEmbeddingModel embeddingModel = new PostgresMlEmbeddingModel(this.jdbcTemplate,
|
||||
PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer("distilbert-base-uncased") // huggingface transformer model name.
|
||||
.withVectorType(VectorType.PG_VECTOR) //vector type in PostgreSQL.
|
||||
.withKwargs(Map.of("device", "cpu")) // optional arguments.
|
||||
.withMetadataMode(MetadataMode.EMBED) // Document metadata mode.
|
||||
.transformer("distilbert-base-uncased") // huggingface transformer model name.
|
||||
.vectorType(VectorType.PG_VECTOR) //vector type in PostgreSQL.
|
||||
.kwargs(Map.of("device", "cpu")) // optional arguments.
|
||||
.metadataMode(MetadataMode.EMBED) // Document metadata mode.
|
||||
.build());
|
||||
|
||||
embeddingModel.afterPropertiesSet(); // initialize the jdbc template and database.
|
||||
|
||||
@@ -48,10 +48,10 @@ public class PostgresMlEmbeddingProperties {
|
||||
|
||||
@NestedConfigurationProperty
|
||||
private PostgresMlEmbeddingOptions options = PostgresMlEmbeddingOptions.builder()
|
||||
.withTransformer(PostgresMlEmbeddingModel.DEFAULT_TRANSFORMER_MODEL)
|
||||
.withVectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
|
||||
.withKwargs(Map.of())
|
||||
.withMetadataMode(MetadataMode.EMBED)
|
||||
.transformer(PostgresMlEmbeddingModel.DEFAULT_TRANSFORMER_MODEL)
|
||||
.vectorType(PostgresMlEmbeddingModel.VectorType.PG_ARRAY)
|
||||
.kwargs(Map.of())
|
||||
.metadataMode(MetadataMode.EMBED)
|
||||
.build();
|
||||
|
||||
public PostgresMlEmbeddingOptions getOptions() {
|
||||
|
||||
Reference in New Issue
Block a user