Register full embedding model names (#654)

Provide backward compatible aliases.
This makes available the same model names that ttok uses.
This commit is contained in:
web-sst 2025-01-22 22:14:03 -06:00 committed by GitHub
parent 59983740e6
commit 6f7ea406bf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 9 additions and 9 deletions

View file

@ -134,14 +134,14 @@ def register_models(register):
@hookimpl
def register_embedding_models(register):
register(
OpenAIEmbeddingModel("ada-002", "text-embedding-ada-002"), aliases=("ada",)
OpenAIEmbeddingModel("text-embedding-ada-002", "text-embedding-ada-002"), aliases=("ada","ada-002",)
)
register(OpenAIEmbeddingModel("3-small", "text-embedding-3-small"))
register(OpenAIEmbeddingModel("3-large", "text-embedding-3-large"))
register(OpenAIEmbeddingModel("text-embedding-3-small", "text-embedding-3-small"), aliases=("3-small",))
register(OpenAIEmbeddingModel("text-embedding-3-large", "text-embedding-3-large"), aliases=("3-large",))
# With varying dimensions
register(OpenAIEmbeddingModel("3-small-512", "text-embedding-3-small", 512))
register(OpenAIEmbeddingModel("3-large-256", "text-embedding-3-large", 256))
register(OpenAIEmbeddingModel("3-large-1024", "text-embedding-3-large", 1024))
register(OpenAIEmbeddingModel("text-embedding-3-small-512", "text-embedding-3-small", 512), aliases=("3-small-512",))
register(OpenAIEmbeddingModel("text-embedding-3-large-256", "text-embedding-3-large", 256), aliases=("3-large-256",))
register(OpenAIEmbeddingModel("text-embedding-3-large-1024", "text-embedding-3-large", 1024), aliases=("3-large-1024",))
class OpenAIEmbeddingModel(EmbeddingModel):

View file

@ -39,7 +39,7 @@ def test_cli_aliases_list(args):
"gpt4 : gpt-4\n"
"4-32k : gpt-4-32k\n"
"e-demo : embed-demo (embedding)\n"
"ada : ada-002 (embedding)\n"
"ada : text-embedding-ada-002 (embedding)\n"
).split("\n"):
line = line.strip()
if not line:
@ -65,7 +65,7 @@ def test_cli_aliases_list_json(args):
"4": "gpt-4",
"gpt4": "gpt-4",
"4-32k": "gpt-4-32k",
"ada": "ada-002",
"ada": "text-embedding-ada-002",
"e-demo": "embed-demo",
}.items()
)

View file

@ -554,7 +554,7 @@ def test_default_embedding_model():
assert result2.exit_code == 0
result3 = runner.invoke(cli, ["embed-models", "default"])
assert result3.exit_code == 0
assert result3.output == "ada-002\n"
assert result3.output == "text-embedding-ada-002\n"
result4 = runner.invoke(cli, ["embed-models", "default", "--remove-default"])
assert result4.exit_code == 0
result5 = runner.invoke(cli, ["embed-models", "default"])