add embedding models

This commit is contained in:
Xi Yan 2025-02-15 14:39:37 -08:00
parent 75e429f950
commit 2e96a9299a
3 changed files with 106 additions and 9 deletions

View file

@ -150,6 +150,36 @@ models:
provider_id: together provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval
model_type: embedding model_type: embedding
- metadata:
embedding_dimension: 768
model_id: togethercomputer/m2-bert-80M-2k-retrieval
provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-2k-retrieval
model_type: embedding
- metadata:
embedding_dimension: 768
model_id: togethercomputer/m2-bert-80M-32k-retrieval
provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-32k-retrieval
model_type: embedding
- metadata:
embedding_dimension: 1024
model_id: WhereIsAI/UAE-Large-V1
provider_id: together
provider_model_id: WhereIsAI/UAE-Large-V1
model_type: embedding
- metadata:
embedding_dimension: 1024
model_id: BAAI/bge-large-en-v1.5
provider_id: together
provider_model_id: BAAI/bge-large-en-v1.5
model_type: embedding
- metadata:
embedding_dimension: 768
model_id: BAAI/bge-base-en-v1.5
provider_id: together
provider_model_id: BAAI/bge-base-en-v1.5
model_type: embedding
shields: shields:
- shield_id: meta-llama/Llama-Guard-3-8B - shield_id: meta-llama/Llama-Guard-3-8B
provider_id: llama-guard provider_id: llama-guard

View file

@ -144,6 +144,36 @@ models:
provider_id: together provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval provider_model_id: togethercomputer/m2-bert-80M-8k-retrieval
model_type: embedding model_type: embedding
- metadata:
embedding_dimension: 768
model_id: togethercomputer/m2-bert-80M-2k-retrieval
provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-2k-retrieval
model_type: embedding
- metadata:
embedding_dimension: 768
model_id: togethercomputer/m2-bert-80M-32k-retrieval
provider_id: together
provider_model_id: togethercomputer/m2-bert-80M-32k-retrieval
model_type: embedding
- metadata:
embedding_dimension: 1024
model_id: WhereIsAI/UAE-Large-V1
provider_id: together
provider_model_id: WhereIsAI/UAE-Large-V1
model_type: embedding
- metadata:
embedding_dimension: 1024
model_id: BAAI/bge-large-en-v1.5
provider_id: together
provider_model_id: BAAI/bge-large-en-v1.5
model_type: embedding
- metadata:
embedding_dimension: 768
model_id: BAAI/bge-base-en-v1.5
provider_id: together
provider_model_id: BAAI/bge-base-en-v1.5
model_type: embedding
shields: shields:
- shield_id: meta-llama/Llama-Guard-3-8B - shield_id: meta-llama/Llama-Guard-3-8B
vector_dbs: [] vector_dbs: []

View file

@ -81,13 +81,50 @@ def get_distribution_template() -> DistributionTemplate:
provider_id="code-interpreter", provider_id="code-interpreter",
), ),
] ]
embedding_model = ModelInput( embedding_models = [
model_id="togethercomputer/m2-bert-80M-8k-retrieval", ModelInput(
provider_id="together", model_id="togethercomputer/m2-bert-80M-8k-retrieval",
model_type=ModelType.embedding, provider_id="together",
provider_model_id="togethercomputer/m2-bert-80M-8k-retrieval", model_type=ModelType.embedding,
metadata={"embedding_dimension": 768}, provider_model_id="togethercomputer/m2-bert-80M-8k-retrieval",
) metadata={"embedding_dimension": 768},
),
ModelInput(
model_id="togethercomputer/m2-bert-80M-2k-retrieval",
provider_id="together",
model_type=ModelType.embedding,
provider_model_id="togethercomputer/m2-bert-80M-2k-retrieval",
metadata={"embedding_dimension": 768},
),
ModelInput(
model_id="togethercomputer/m2-bert-80M-32k-retrieval",
provider_id="together",
model_type=ModelType.embedding,
provider_model_id="togethercomputer/m2-bert-80M-32k-retrieval",
metadata={"embedding_dimension": 768},
),
ModelInput(
model_id="WhereIsAI/UAE-Large-V1",
provider_id="together",
model_type=ModelType.embedding,
provider_model_id="WhereIsAI/UAE-Large-V1",
metadata={"embedding_dimension": 1024},
),
ModelInput(
model_id="BAAI/bge-large-en-v1.5",
provider_id="together",
model_type=ModelType.embedding,
provider_model_id="BAAI/bge-large-en-v1.5",
metadata={"embedding_dimension": 1024},
),
ModelInput(
model_id="BAAI/bge-base-en-v1.5",
provider_id="together",
model_type=ModelType.embedding,
provider_model_id="BAAI/bge-base-en-v1.5",
metadata={"embedding_dimension": 768},
),
]
return DistributionTemplate( return DistributionTemplate(
name=name, name=name,
@ -103,7 +140,7 @@ def get_distribution_template() -> DistributionTemplate:
"inference": [inference_provider, embedding_provider], "inference": [inference_provider, embedding_provider],
"vector_io": [vector_io_provider], "vector_io": [vector_io_provider],
}, },
default_models=default_models + [embedding_model], default_models=default_models + embedding_models,
default_tool_groups=default_tool_groups, default_tool_groups=default_tool_groups,
default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")], default_shields=[ShieldInput(shield_id="meta-llama/Llama-Guard-3-8B")],
), ),
@ -134,7 +171,7 @@ def get_distribution_template() -> DistributionTemplate:
}, },
default_models=[ default_models=[
*default_models, *default_models,
embedding_model, *embedding_models,
], ],
default_shields=[ default_shields=[
ShieldInput( ShieldInput(