Hopiu/llm
1
0
Fork 0
mirror of https://github.com/Hopiu/llm.git synced 2026-03-17 05:00:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add embedding model aliases to llm aliases, refs #192

Browse source
This commit is contained in:
Simon Willison 2023-09-03 17:55:56 -07:00
parent 408297f1f5
commit 2e90a30c4f
3 changed files with 23 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/aliases.md
View file

@ -26,6 +26,7 @@ chatgpt-16k : gpt-3.5-turbo-16k
4 : gpt-4
gpt4 : gpt-4
4-32k : gpt-4-32k
ada : ada-002 (embedding)
```
<!-- [[[end]]] -->

22
llm/cli.py
View file

@ -11,6 +11,7 @@ from llm import (
UnknownModelError,
encode,
get_embedding_models_with_aliases,
get_embedding_model_aliases,
get_embedding_model,
get_key,
get_plugins,
@ -735,14 +736,23 @@ def aliases_list(json_):
to_output = []
for alias, model in get_model_aliases().items():
if alias != model.model_id:
to_output.append((alias, model.model_id))
to_output.append((alias, model.model_id, ""))
for alias, embedding_model in get_embedding_model_aliases().items():
if alias != embedding_model.model_id:
to_output.append((alias, embedding_model.model_id, "embedding"))
if json_:
click.echo(json.dumps({key: value for key, value in to_output}, indent=4))
click.echo(
json.dumps({key: value for key, value, type_ in to_output}, indent=4)
)
return
max_alias_length = max(len(a) for a, _ in to_output)
fmt = "{alias:<" + str(max_alias_length) + "} : {model_id}"
for alias, model_id in to_output:
click.echo(fmt.format(alias=alias, model_id=model_id))
max_alias_length = max(len(a) for a, _, _ in to_output)
fmt = "{alias:<" + str(max_alias_length) + "} : {model_id}{type_}"
for alias, model_id, type_ in to_output:
click.echo(
fmt.format(
alias=alias, model_id=model_id, type_=f" ({type_})" if type_ else ""
)
)
@aliases.command(name="set")

6
tests/test_aliases.py
View file

@ -25,6 +25,7 @@ def test_remove_alias():
@pytest.mark.parametrize("args", (["aliases", "list"], ["aliases"]))
def test_cli_aliases_list(args):
llm.set_alias("e-demo", "embed-demo")
runner = CliRunner()
result = runner.invoke(cli, args)
assert result.exit_code == 0
@ -36,11 +37,14 @@ def test_cli_aliases_list(args):
"4 : gpt-4\n"
"gpt4 : gpt-4\n"
"4-32k : gpt-4-32k\n"
"e-demo : embed-demo (embedding)\n"
"ada : ada-002 (embedding)\n"
)
@pytest.mark.parametrize("args", (["aliases", "list"], ["aliases"]))
def test_cli_aliases_list_json(args):
llm.set_alias("e-demo", "embed-demo")
runner = CliRunner()
result = runner.invoke(cli, args + ["--json"])
assert result.exit_code == 0
@ -52,6 +56,8 @@ def test_cli_aliases_list_json(args):
"4": "gpt-4",
"gpt4": "gpt-4",
"4-32k": "gpt-4-32k",
"ada": "ada-002",
"e-demo": "embed-demo",
}