diff options
Diffstat (limited to 'g4f/locals/models.py')
| -rw-r--r-- | g4f/locals/models.py | 51 |
1 files changed, 51 insertions, 0 deletions
diff --git a/g4f/locals/models.py b/g4f/locals/models.py new file mode 100644 index 00000000..f82f5448 --- /dev/null +++ b/g4f/locals/models.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import os +import requests +import json + +from ..requests.raise_for_status import raise_for_status + +def load_models(): + response = requests.get("https://gpt4all.io/models/models3.json") + raise_for_status(response) + return format_models(response.json()) + +def get_model_name(filename: str) -> str: + name = filename.split(".", 1)[0] + for replace in ["-v1_5", "-v1", "-q4_0", "_v01", "-v0", "-f16", "-gguf2", "-newbpe"]: + name = name.replace(replace, "") + return name + +def format_models(models: list) -> dict: + return {get_model_name(model["filename"]): { + "path": model["filename"], + "ram": model["ramrequired"], + "prompt": model["promptTemplate"] if "promptTemplate" in model else None, + "system": model["systemPrompt"] if "systemPrompt" in model else None, + } for model in models} + +def read_models(file_path: str): + with open(file_path, "rb") as f: + return json.load(f) + +def save_models(file_path: str, data): + with open(file_path, 'w') as f: + json.dump(data, f, indent=4) + +def get_model_dir() -> str: + local_dir = os.path.dirname(os.path.abspath(__file__)) + project_dir = os.path.dirname(os.path.dirname(local_dir)) + model_dir = os.path.join(project_dir, "models") + if os.path.exists(model_dir): + return model_dir + +def get_models() -> dict[str, dict]: + model_dir = get_model_dir() + file_path = os.path.join(model_dir, "models.json") + if os.path.isfile(file_path): + return read_models(file_path) + else: + models = load_models() + save_models(file_path, models) + return models
\ No newline at end of file |