|
21 | 21 | import org.elasticsearch.xpack.inference.services.voyageai.action.VoyageAIActionVisitor; |
22 | 22 |
|
23 | 23 | import java.net.URI; |
| 24 | +import java.util.Collections; |
| 25 | +import java.util.HashMap; |
24 | 26 | import java.util.Map; |
25 | 27 | import java.util.Objects; |
26 | 28 |
|
27 | 29 | public abstract class VoyageAIModel extends RateLimitGroupingModel { |
28 | 30 | private static final String DEFAULT_MODEL_FAMILY = "default_model_family"; |
29 | | - private static final Map<String, String> MODEL_TO_MODEL_FAMILY = Map.of( |
30 | | - "voyage-multimodal-3", |
31 | | - "embed_multimodal", |
32 | | - "voyage-3-large", |
33 | | - "embed_large", |
34 | | - "voyage-code-3", |
35 | | - "embed_large", |
36 | | - "voyage-3", |
37 | | - "embed_medium", |
38 | | - "voyage-3-lite", |
39 | | - "embed_small", |
40 | | - "voyage-finance-2", |
41 | | - "embed_large", |
42 | | - "voyage-law-2", |
43 | | - "embed_large", |
44 | | - "voyage-code-2", |
45 | | - "embed_large", |
46 | | - "rerank-2", |
47 | | - "rerank_large", |
48 | | - "rerank-2-lite", |
49 | | - "rerank_small" |
50 | | - ); |
| 31 | + private static final Map<String, String> MODEL_TO_MODEL_FAMILY; |
| 32 | + |
| 33 | + static { |
| 34 | + Map<String, String> tempMap = new HashMap<>(); |
| 35 | + tempMap.put("voyage-3.5", "embed_medium"); |
| 36 | + tempMap.put("voyage-3.5-lite", "embed_small"); |
| 37 | + tempMap.put("voyage-multimodal-3", "embed_multimodal"); |
| 38 | + tempMap.put("voyage-3-large", "embed_large"); |
| 39 | + tempMap.put("voyage-code-3", "embed_large"); |
| 40 | + tempMap.put("voyage-3", "embed_medium"); |
| 41 | + tempMap.put("voyage-3-lite", "embed_small"); |
| 42 | + tempMap.put("voyage-finance-2", "embed_large"); |
| 43 | + tempMap.put("voyage-law-2", "embed_large"); |
| 44 | + tempMap.put("voyage-code-2", "embed_large"); |
| 45 | + tempMap.put("rerank-2", "rerank_large"); |
| 46 | + tempMap.put("rerank-2-lite", "rerank_small"); |
| 47 | + |
| 48 | + MODEL_TO_MODEL_FAMILY = Collections.unmodifiableMap(tempMap); |
| 49 | + } |
51 | 50 |
|
52 | 51 | private final SecureString apiKey; |
53 | 52 | private final VoyageAIRateLimitServiceSettings rateLimitServiceSettings; |
|
0 commit comments