File size: 1,961 Bytes
836d634
53bd9f8
 
e07a1f9
 
 
b059201
 
 
53bd9f8
836d634
 
 
 
e07a1f9
abab5af
 
836d634
 
 
b059201
 
 
 
 
 
 
 
abab5af
e07a1f9
 
 
 
 
b059201
e07a1f9
 
b059201
 
 
836d634
e07a1f9
 
 
abab5af
 
53bd9f8
e07a1f9
836d634
 
 
 
 
 
e07a1f9
836d634
 
 
 
 
 
 
 
53bd9f8
e07a1f9
 
 
 
 
b059201
 
 
 
 
 
836d634
 
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
from flask import Flask, request, jsonify
import subprocess

MODEL_PATH = "/app/model/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
BINARY_PATH = "/llama.cpp/build/bin/llama-cli"

MAX_TOKENS = "128"
THREADS = "2"
CONTEXT_SIZE = "2048"

app = Flask(__name__)

def query_model(system: str, user: str) -> str:
    prompt = f"System: {system}\nUser: {user}" if system else user

    try:
        cmd = [
            BINARY_PATH,
            "-m", MODEL_PATH,
            "-t", THREADS,
            "-c", CONTEXT_SIZE,
            "-p", prompt,
            "--n_predict", MAX_TOKENS,
            "--temp", "0.7",
            "--top_k", "40",
            "--top_p", "0.9",
            "--repeat_penalty", "1.1",
            "--no-warmup"
        ]

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            timeout=90  # allow slightly longer for HF
        )

        if result.returncode != 0:
            return f"Model error:\n{result.stderr.strip()}"

        return result.stdout.strip() or "No output from model."

    except subprocess.TimeoutExpired:
        return "Error: Model took too long (timeout)."
    except Exception as e:
        return f"Unexpected error:\n{str(e)}"


@app.route("/api/query", methods=["GET", "POST"])
def api_query():
    if request.method == "POST":
        data = request.json or {}
        user_prompt = data.get("prompt", "")
        system_prompt = data.get("system", "")
    else:
        user_prompt = request.args.get("prompt", "")
        system_prompt = request.args.get("system", "")

    if not user_prompt.strip():
        return jsonify({"error": "Prompt is required"}), 400

    output = query_model(system_prompt, user_prompt)
    return jsonify({"output": output})


@app.route("/")
def home():
    return "API is running"


@app.route("/health")
def health():
    return {"status": "ok"}


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)