feat: Add support for yaml based configs

abetlen · abetlen · commit 060bfa64d529 · 2024-04-10T02:47:01.000-04:00
diff --git a/examples/batch-processing/server.py b/examples/batch-processing/server.py
@@ -0,0 +1,30 @@
+"""llama-cpp-python server from scratch in a single file.
+"""
+
+# import llama_cpp
+
+# path = b"../../models/Qwen1.5-0.5B-Chat-GGUF/qwen1_5-0_5b-chat-q8_0.gguf"
+
+# model_params = llama_cpp.llama_model_default_params()
+# model = llama_cpp.llama_load_model_from_file(path, model_params)
+
+# if model is None:
+#     raise RuntimeError(f"Failed to load model from file: {path}")
+
+
+# ctx_params = llama_cpp.llama_context_default_params()
+# ctx = llama_cpp.llama_new_context_with_model(model, ctx_params)
+
+# if ctx is None:
+#     raise RuntimeError("Failed to create context")
+
+
+from fastapi import FastAPI
+
+app = FastAPI()
+
+import openai.types.chat as types
+
+@app.post("/v1/chat/completions")
+def create_chat_completions():
+    return {"message": "Hello World"}
diff --git a/llama_cpp/server/__main__.py b/llama_cpp/server/__main__.py
@@ -59,7 +59,16 @@ def main():
             if not os.path.exists(config_file):
                 raise ValueError(f"Config file {config_file} not found!")
             with open(config_file, "rb") as f:
-                config_file_settings = ConfigFileSettings.model_validate_json(f.read())
+                # Check if yaml file
+                if config_file.endswith(".yaml") or config_file.endswith(".yml"):
+                    import yaml
+                    import json
+
+                    config_file_settings = ConfigFileSettings.model_validate_json(
+                        json.dumps(yaml.safe_load(f))
+                    )
+                else:
+                    config_file_settings = ConfigFileSettings.model_validate_json(f.read())
                 server_settings = ServerSettings.model_validate(config_file_settings)
                 model_settings = config_file_settings.models
         else:
diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py
@@ -97,7 +97,15 @@ def create_app(
         if not os.path.exists(config_file):
             raise ValueError(f"Config file {config_file} not found!")
         with open(config_file, "rb") as f:
-            config_file_settings = ConfigFileSettings.model_validate_json(f.read())
+            # Check if yaml file
+            if config_file.endswith(".yaml") or config_file.endswith(".yml"):
+                import yaml
+
+                config_file_settings = ConfigFileSettings.model_validate_json(
+                    json.dumps(yaml.safe_load(f))
+                )
+            else:
+                config_file_settings = ConfigFileSettings.model_validate_json(f.read())
             server_settings = ServerSettings.model_validate(config_file_settings)
             model_settings = config_file_settings.models
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -35,6 +35,7 @@ server = [
     "pydantic-settings>=2.0.1",
     "sse-starlette>=1.6.1",
     "starlette-context>=0.3.6,<0.4",
+    "PyYAML>=5.1",
 ]
 test = [
     "pytest>=7.4.0",

Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,7 @@ server = [`
`35`	`35`	`"pydantic-settings>=2.0.1",`
`36`	`36`	`"sse-starlette>=1.6.1",`
`37`	`37`	`"starlette-context>=0.3.6,<0.4",`
	`38`	`+ "PyYAML>=5.1",`
`38`	`39`	`]`
`39`	`40`	`test = [`
`40`	`41`	`"pytest>=7.4.0",`