forked from abetlen/llama-cpp-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathcli.py
More file actions
129 lines (103 loc) · 4.09 KB
/
cli.py
File metadata and controls
129 lines (103 loc) · 4.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
from __future__ import annotations
import argparse
import json
from typing import List, Literal, Union, Any, Type, TypeVar, Dict
from pydantic import BaseModel
def _get_base_type(annotation: Type[Any]) -> Type[Any]:
if getattr(annotation, "__origin__", None) is Literal:
assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore
return type(annotation.__args__[0]) # type: ignore
elif getattr(annotation, "__origin__", None) is Union:
assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore
non_optional_args: List[Type[Any]] = [
arg
for arg in annotation.__args__
if arg is not type(None) # type: ignore
]
if non_optional_args:
return _get_base_type(non_optional_args[0])
elif (
getattr(annotation, "__origin__", None) is list
or getattr(annotation, "__origin__", None) is List
):
assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore
return _get_base_type(annotation.__args__[0]) # type: ignore
return annotation
def _contains_list_type(annotation: Type[Any] | None) -> bool:
origin = getattr(annotation, "__origin__", None)
if origin is list or origin is List:
return True
elif origin in (Literal, Union):
return any(_contains_list_type(arg) for arg in annotation.__args__) # type: ignore
else:
return False
def _contains_dict_type(annotation: Type[Any] | None) -> bool:
origin = getattr(annotation, "__origin__", None)
if origin is dict or origin is Dict:
return True
elif origin in (Literal, Union):
return any(_contains_dict_type(arg) for arg in annotation.__args__) # type: ignore
else:
return False
def _parse_bool_arg(arg: str | bytes | bool) -> bool:
if isinstance(arg, bytes):
arg = arg.decode("utf-8")
true_values = {"1", "on", "t", "true", "y", "yes"}
false_values = {"0", "off", "f", "false", "n", "no"}
arg_str = str(arg).lower().strip()
if arg_str in true_values:
return True
elif arg_str in false_values:
return False
else:
raise ValueError(f"Invalid boolean argument: {arg}")
def _parse_json_object_arg(arg: str | bytes) -> dict[str, Any]:
if isinstance(arg, bytes):
arg = arg.decode("utf-8")
value = json.loads(arg)
if not isinstance(value, dict):
raise ValueError(f"Invalid JSON object argument: {arg}")
return value
def add_args_from_model(parser: argparse.ArgumentParser, model: Type[BaseModel]):
"""Add arguments from a pydantic model to an argparse parser."""
for name, field in model.model_fields.items():
description = field.description
if field.default and description and not field.is_required():
description += f" (default: {field.default})"
base_type = (
_get_base_type(field.annotation) if field.annotation is not None else str
)
list_type = _contains_list_type(field.annotation)
dict_type = _contains_dict_type(field.annotation)
if dict_type:
parser.add_argument(
f"--{name}",
dest=name,
type=_parse_json_object_arg,
help=description,
)
elif base_type is not bool:
parser.add_argument(
f"--{name}",
dest=name,
nargs="*" if list_type else None,
type=base_type,
help=description,
)
if base_type is bool:
parser.add_argument(
f"--{name}",
dest=name,
type=_parse_bool_arg,
help=f"{description}",
)
T = TypeVar("T", bound=Type[BaseModel])
def parse_model_from_args(model: T, args: argparse.Namespace) -> T:
"""Parse a pydantic model from an argparse namespace."""
return model(
**{
k: v
for k, v in vars(args).items()
if v is not None and k in model.model_fields
}
)