Skip to content

Commit 525417a

Browse files
kderussoJeremyDahlgren
authored andcommitted
Remove upper limit for chunking settings (elastic#133718)
1 parent 3863edc commit 525417a

File tree

14 files changed

+193
-21
lines changed

14 files changed

+193
-21
lines changed

docs/changelog/133718.yaml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 133718
2+
summary: Remove upper limit for chunking settings
3+
area: Machine Learning
4+
type: enhancement
5+
issues: []

server/src/main/java/org/elasticsearch/inference/ChunkingSettings.java

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,13 @@
1515
import java.util.Map;
1616

1717
public interface ChunkingSettings extends ToXContentObject, VersionedNamedWriteable {
18+
1819
ChunkingStrategy getChunkingStrategy();
1920

2021
Map<String, Object> asMap();
22+
23+
/**
24+
* @return The max chunk size specified, or null if not specified
25+
*/
26+
Integer maxChunkSize();
2127
}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/NoneChunkingSettings.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,11 @@ public ChunkingStrategy getChunkingStrategy() {
3737
return STRATEGY;
3838
}
3939

40+
@Override
41+
public Integer maxChunkSize() {
42+
return null;
43+
}
44+
4045
@Override
4146
public String getWriteableName() {
4247
return NAME;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/RecursiveChunker.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ public List<ChunkOffset> chunk(String input, ChunkingSettings chunkingSettings)
4141
input,
4242
new ChunkOffset(0, input.length()),
4343
recursiveChunkingSettings.getSeparators(),
44-
recursiveChunkingSettings.getMaxChunkSize(),
44+
recursiveChunkingSettings.maxChunkSize(),
4545
0
4646
);
4747
} else {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/RecursiveChunkingSettings.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ public class RecursiveChunkingSettings implements ChunkingSettings {
3131
public static final String NAME = "RecursiveChunkingSettings";
3232
private static final ChunkingStrategy STRATEGY = ChunkingStrategy.RECURSIVE;
3333
private static final int MAX_CHUNK_SIZE_LOWER_LIMIT = 10;
34-
private static final int MAX_CHUNK_SIZE_UPPER_LIMIT = 300;
3534

3635
private static final Set<String> VALID_KEYS = Set.of(
3736
ChunkingSettingsOptions.STRATEGY.toString(),
@@ -63,11 +62,10 @@ public static RecursiveChunkingSettings fromMap(Map<String, Object> map) {
6362
);
6463
}
6564

66-
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerBetween(
65+
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerGreaterThanOrEqualToMin(
6766
map,
6867
ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
6968
MAX_CHUNK_SIZE_LOWER_LIMIT,
70-
MAX_CHUNK_SIZE_UPPER_LIMIT,
7169
ModelConfigurations.CHUNKING_SETTINGS,
7270
validationException
7371
);
@@ -105,7 +103,8 @@ public static RecursiveChunkingSettings fromMap(Map<String, Object> map) {
105103
return new RecursiveChunkingSettings(maxChunkSize, separators);
106104
}
107105

108-
public int getMaxChunkSize() {
106+
@Override
107+
public Integer maxChunkSize() {
109108
return maxChunkSize;
110109
}
111110

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/SentenceBoundaryChunkingSettings.java

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,6 @@ public class SentenceBoundaryChunkingSettings implements ChunkingSettings {
3131
public static final String NAME = "SentenceBoundaryChunkingSettings";
3232
private static final ChunkingStrategy STRATEGY = ChunkingStrategy.SENTENCE;
3333
private static final int MAX_CHUNK_SIZE_LOWER_LIMIT = 20;
34-
private static final int MAX_CHUNK_SIZE_UPPER_LIMIT = 300;
3534
private static final Set<String> VALID_KEYS = Set.of(
3635
ChunkingSettingsOptions.STRATEGY.toString(),
3736
ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
@@ -55,6 +54,11 @@ public SentenceBoundaryChunkingSettings(StreamInput in) throws IOException {
5554
}
5655
}
5756

57+
@Override
58+
public Integer maxChunkSize() {
59+
return maxChunkSize;
60+
}
61+
5862
@Override
5963
public Map<String, Object> asMap() {
6064
return Map.of(
@@ -77,11 +81,10 @@ public static SentenceBoundaryChunkingSettings fromMap(Map<String, Object> map)
7781
);
7882
}
7983

80-
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerBetween(
84+
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerGreaterThanOrEqualToMin(
8185
map,
8286
ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
8387
MAX_CHUNK_SIZE_LOWER_LIMIT,
84-
MAX_CHUNK_SIZE_UPPER_LIMIT,
8588
ModelConfigurations.CHUNKING_SETTINGS,
8689
validationException
8790
);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/chunking/WordBoundaryChunkingSettings.java

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ public class WordBoundaryChunkingSettings implements ChunkingSettings {
3030
public static final String NAME = "WordBoundaryChunkingSettings";
3131
private static final ChunkingStrategy STRATEGY = ChunkingStrategy.WORD;
3232
private static final int MAX_CHUNK_SIZE_LOWER_LIMIT = 10;
33-
private static final int MAX_CHUNK_SIZE_UPPER_LIMIT = 300;
3433
private static final Set<String> VALID_KEYS = Set.of(
3534
ChunkingSettingsOptions.STRATEGY.toString(),
3635
ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
@@ -61,7 +60,8 @@ public Map<String, Object> asMap() {
6160
);
6261
}
6362

64-
public int maxChunkSize() {
63+
@Override
64+
public Integer maxChunkSize() {
6565
return maxChunkSize;
6666
}
6767

@@ -79,11 +79,10 @@ public static WordBoundaryChunkingSettings fromMap(Map<String, Object> map) {
7979
);
8080
}
8181

82-
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerBetween(
82+
Integer maxChunkSize = ServiceUtils.extractRequiredPositiveIntegerGreaterThanOrEqualToMin(
8383
map,
8484
ChunkingSettingsOptions.MAX_CHUNK_SIZE.toString(),
8585
MAX_CHUNK_SIZE_LOWER_LIMIT,
86-
MAX_CHUNK_SIZE_UPPER_LIMIT,
8786
ModelConfigurations.CHUNKING_SETTINGS,
8887
validationException
8988
);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ServiceUtils.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -749,6 +749,25 @@ public static Integer extractRequiredPositiveIntegerLessThanOrEqualToMax(
749749
return field;
750750
}
751751

752+
public static Integer extractRequiredPositiveIntegerGreaterThanOrEqualToMin(
753+
Map<String, Object> map,
754+
String settingName,
755+
int minValue,
756+
String scope,
757+
ValidationException validationException
758+
) {
759+
Integer field = extractRequiredPositiveInteger(map, settingName, scope, validationException);
760+
761+
if (field != null && field < minValue) {
762+
validationException.addValidationError(
763+
ServiceUtils.mustBeGreaterThanOrEqualNumberErrorMessage(settingName, scope, field, minValue)
764+
);
765+
return null;
766+
}
767+
768+
return field;
769+
}
770+
752771
public static Integer extractRequiredPositiveIntegerBetween(
753772
Map<String, Object> map,
754773
String settingName,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/ElserInternalModel.java

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212

1313
public class ElserInternalModel extends ElasticsearchInternalModel {
1414

15+
// Ensure that inference endpoints based on ELSER don't go past its truncation window of 512 tokens
16+
public static final int ELSER_MAX_WINDOW_SIZE = 300;
17+
1518
public ElserInternalModel(
1619
String inferenceEntityId,
1720
TaskType taskType,
@@ -21,6 +24,14 @@ public ElserInternalModel(
2124
ChunkingSettings chunkingSettings
2225
) {
2326
super(inferenceEntityId, taskType, service, serviceSettings, taskSettings, chunkingSettings);
27+
if (chunkingSettings != null && chunkingSettings.maxChunkSize() != null) {
28+
if (chunkingSettings.maxChunkSize() > ELSER_MAX_WINDOW_SIZE) throw new IllegalArgumentException(
29+
"ELSER based models do not support chunk sizes larger than "
30+
+ ELSER_MAX_WINDOW_SIZE
31+
+ ". Requested chunk size: "
32+
+ chunkingSettings.maxChunkSize()
33+
);
34+
}
2435
}
2536

2637
@Override

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elasticsearch/MultilingualE5SmallModel.java

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212

1313
public class MultilingualE5SmallModel extends ElasticsearchInternalModel {
1414

15+
// Ensure that inference endpoints based on E5 small don't go past its window size
16+
public static final int E5_SMALL_MAX_WINDOW_SIZE = 300;
17+
1518
public MultilingualE5SmallModel(
1619
String inferenceEntityId,
1720
TaskType taskType,
@@ -20,6 +23,15 @@ public MultilingualE5SmallModel(
2023
ChunkingSettings chunkingSettings
2124
) {
2225
super(inferenceEntityId, taskType, service, serviceSettings, chunkingSettings);
26+
if (chunkingSettings != null && chunkingSettings.maxChunkSize() != null) {
27+
if (chunkingSettings.maxChunkSize() > E5_SMALL_MAX_WINDOW_SIZE) throw new IllegalArgumentException(
28+
serviceSettings.modelId()
29+
+ " does not support chunk sizes larger than "
30+
+ E5_SMALL_MAX_WINDOW_SIZE
31+
+ ". Requested chunk size: "
32+
+ chunkingSettings.maxChunkSize()
33+
);
34+
}
2335
}
2436

2537
@Override

0 commit comments

Comments
 (0)