Update server.py

server.py

@@ -1,102 +1,69 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import hf_hub_download
-import requests
 import subprocess
+import requests
 import uvicorn
 import os
-import json
+import time
+import hashlib
 
 app = FastAPI()
 
-MODELS = {
-    "tinyllama": "models/tinyllama.gguf"
-}
-
-class ChatRequest(BaseModel):
-    model: str
-    messages: list
-
-class GenerateRequest(BaseModel):
-    model: str
-    prompt: str
-
-
 # ---------------------------
-# logging
+# MODEL CONFIG
 # ---------------------------
 
-def log(title, data):
-    print("\n==============================")
-    print(title)
-    print(data)
-    print("==============================\n")
+MODEL_NAME = "tinyllama"
+MODEL_PATH = "models/tinyllama.gguf"
+
+MODEL_REPO = "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF"
+MODEL_FILE = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
 
 
 # ---------------------------
-# prompt builder
+# REQUEST MODELS
 # ---------------------------
 
-def build_prompt(messages):
-
-    prompt = ""
-
-    for m in messages:
-        role = m.get("role", "user")
-        content = m.get("content", "")
-
-        if content.strip() == "":
-            continue
-
-        prompt += f"{role}: {content}\n"
-
-    prompt += "assistant:"
+class ChatRequest(BaseModel):
+    model: str
+    messages: list
 
-    log("PROMPT", prompt)
 
-    return prompt
+class GenerateRequest(BaseModel):
+    model: str
+    prompt: str
 
 
 # ---------------------------
-# download model
+# DOWNLOAD MODEL
 # ---------------------------
 
 os.makedirs("models", exist_ok=True)
 
-MODEL_FILES = {
-    "tinyllama": (
-        "TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF",
-        "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf"
-    )
-}
-
-for name, (repo, file) in MODEL_FILES.items():
-
-    path = f"models/{name}.gguf"
+if not os.path.exists(MODEL_PATH):
 
-    if not os.path.exists(path):
+    print("Downloading model from HuggingFace")
 
-        print(f"Downloading model {name}")
-
-        downloaded = hf_hub_download(
-            repo_id=repo,
-            filename=file
-        )
+    downloaded = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILE
+    )
 
-        os.system(f"cp {downloaded} {path}")
+    os.system(f"cp {downloaded} {MODEL_PATH}")
 
-        print(f"Model ready: {path}")
+    print("Model downloaded")
 
 
 # ---------------------------
-# start llama-server
+# START LLAMA SERVER
 # ---------------------------
 
 print("Starting llama-server...")
 
 subprocess.Popen([
     "./llama.cpp/build/bin/llama-server",
-    "-m", "models/tinyllama.gguf",
+    "-m", MODEL_PATH,
     "--host", "0.0.0.0",
     "--port", "8080",
     "-c", "2048"
@@ -104,7 +71,30 @@ subprocess.Popen([
 
 
 # ---------------------------
-# root
+# PROMPT BUILDER
+# ---------------------------
+
+def build_prompt(messages):
+
+    prompt = ""
+
+    for m in messages:
+
+        role = m.get("role")
+        content = m.get("content", "").strip()
+
+        if content == "":
+            continue
+
+        prompt += f"{role}: {content}\n"
+
+    prompt += "assistant:"
+
+    return prompt
+
+
+# ---------------------------
+# ROOT
 # ---------------------------
 
 @app.get("/")
@@ -113,74 +103,119 @@ def root():
 
 
 # ---------------------------
-# model list
+# MODEL LIST (OLLAMA FORMAT)
 # ---------------------------
 
 @app.get("/api/tags")
-def list_models():
+def tags():
+
+    size = os.path.getsize(MODEL_PATH)
+
+    digest = hashlib.sha256(open(MODEL_PATH, "rb").read()).hexdigest()
 
     return {
         "models": [
-            {"name": "tinyllama"}
+            {
+                "name": MODEL_NAME,
+                "model": MODEL_NAME,
+                "modified_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+                "size": size,
+                "digest": digest,
+                "details": {
+                    "format": "gguf",
+                    "family": "llama",
+                    "families": ["llama"],
+                    "parameter_size": "1.1B",
+                    "quantization_level": "Q4_K_M"
+                }
+            }
         ]
     }
 
 
 # ---------------------------
-# chat endpoint
+# GENERATE ENDPOINT
 # ---------------------------
 
-@app.post("/api/chat")
-def chat(req: ChatRequest):
+@app.post("/api/generate")
+def generate(req: GenerateRequest):
 
-    prompt = build_prompt(req.messages)
+    start = time.time()
 
     response = requests.post(
         "http://localhost:8080/completion",
         json={
-            "prompt": prompt,
+            "prompt": req.prompt,
             "n_predict": 200
         }
     )
 
-    data = response.json()
+    text = response.json()["content"].strip()
+
+    duration = int((time.time() - start) * 1e9)
 
     return {
         "model": req.model,
-        "message": {
-            "role": "assistant",
-            "content": data["content"]
-        },
-        "done": True
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "response": text,
+        "done": True,
+        "done_reason": "stop",
+        "total_duration": duration,
+        "load_duration": 0,
+        "prompt_eval_count": 0,
+        "prompt_eval_duration": 0,
+        "eval_count": len(text.split()),
+        "eval_duration": duration
     }
 
 
 # ---------------------------
-# generate endpoint
+# CHAT ENDPOINT
 # ---------------------------
 
-@app.post("/api/generate")
-def generate(req: GenerateRequest):
+@app.post("/api/chat")
+def chat(req: ChatRequest):
+
+    start = time.time()
+
+    prompt = build_prompt(req.messages)
 
     response = requests.post(
         "http://localhost:8080/completion",
         json={
-            "prompt": req.prompt,
+            "prompt": prompt,
             "n_predict": 200
         }
     )
 
-    data = response.json()
+    text = response.json()["content"].strip()
+
+    duration = int((time.time() - start) * 1e9)
 
     return {
         "model": req.model,
-        "response": data["content"],
-        "done": True
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "message": {
+            "role": "assistant",
+            "content": text,
+            "thinking": "",
+            "tool_calls": [],
+            "images": []
+        },
+        "done": True,
+        "done_reason": "stop",
+        "total_duration": duration,
+        "load_duration": 0,
+        "prompt_eval_count": 0,
+        "prompt_eval_duration": 0,
+        "eval_count": len(text.split()),
+        "eval_duration": duration,
+        "logprobs": []
     }
 
 
 # ---------------------------
-# start API
+# START API
 # ---------------------------
 
 if __name__ == "__main__":