add embedding model by default to distribution templates (#617)

# What does this PR do?
Adds the sentence transformer provider and the `all-MiniLM-L6-v2`
embedding model to the default models to register in the run.yaml for
all providers.

## Test Plan
llama stack build --template together --image-type conda
llama stack run
~/.llama/distributions/llamastack-together/together-run.yaml
This commit is contained in:
Dinesh Yeduguru 2024-12-13 12:48:00 -08:00 committed by GitHub
parent e893b22868
commit 516e1a3e59
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 473 additions and 64 deletions

View file

@ -79,10 +79,12 @@ models:
model_id: ${env.INFERENCE_MODEL}
provider_id: tgi-inference
provider_model_id: null
model_type: llm
- metadata: {}
model_id: ${env.SAFETY_MODEL}
provider_id: tgi-safety
provider_model_id: null
model_type: llm
shields:
- params: null
shield_id: ${env.SAFETY_MODEL}

View file

@ -17,6 +17,9 @@ providers:
provider_type: remote::tgi
config:
url: ${env.TGI_URL}
- provider_id: sentence-transformers
provider_type: inline::sentence-transformers
config: {}
memory:
- provider_id: faiss
provider_type: inline::faiss
@ -75,6 +78,13 @@ models:
model_id: ${env.INFERENCE_MODEL}
provider_id: tgi-inference
provider_model_id: null
model_type: llm
- metadata:
embedding_dimension: 384
model_id: all-MiniLM-L6-v2
provider_id: sentence-transformers
provider_model_id: null
model_type: embedding
shields: []
memory_banks: []
datasets: []

View file

@ -6,7 +6,12 @@
from pathlib import Path
from llama_stack.apis.models.models import ModelType
from llama_stack.distribution.datatypes import ModelInput, Provider, ShieldInput
from llama_stack.providers.inline.inference.sentence_transformers import (
SentenceTransformersInferenceConfig,
)
from llama_stack.providers.inline.memory.faiss.config import FaissImplConfig
from llama_stack.providers.remote.inference.tgi import TGIImplConfig
from llama_stack.templates.template import DistributionTemplate, RunConfigSettings
@ -31,6 +36,11 @@ def get_distribution_template() -> DistributionTemplate:
url="${env.TGI_URL}",
),
)
embedding_provider = Provider(
provider_id="sentence-transformers",
provider_type="inline::sentence-transformers",
config=SentenceTransformersInferenceConfig.sample_run_config(),
)
memory_provider = Provider(
provider_id="faiss",
provider_type="inline::faiss",
@ -41,6 +51,14 @@ def get_distribution_template() -> DistributionTemplate:
model_id="${env.INFERENCE_MODEL}",
provider_id="tgi-inference",
)
embedding_model = ModelInput(
model_id="all-MiniLM-L6-v2",
provider_id="sentence-transformers",
model_type=ModelType.embedding,
metadata={
"embedding_dimension": 384,
},
)
safety_model = ModelInput(
model_id="${env.SAFETY_MODEL}",
provider_id="tgi-safety",
@ -57,10 +75,10 @@ def get_distribution_template() -> DistributionTemplate:
run_configs={
"run.yaml": RunConfigSettings(
provider_overrides={
"inference": [inference_provider],
"inference": [inference_provider, embedding_provider],
"memory": [memory_provider],
},
default_models=[inference_model],
default_models=[inference_model, embedding_model],
),
"run-with-safety.yaml": RunConfigSettings(
provider_overrides={