File size: 3,734 Bytes
8e4b2e1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
import os

# Set these before importing llama_cpp.
# They affect BLAS/OpenMP-style CPU threading.
CPU_COUNT = os.cpu_count() or 2
CPU_THREADS = int(os.getenv("CPU_THREADS", str(CPU_COUNT)))

os.environ.setdefault("OMP_NUM_THREADS", str(CPU_THREADS))
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(CPU_THREADS))
os.environ.setdefault("MKL_NUM_THREADS", str(CPU_THREADS))
os.environ.setdefault("NUMEXPR_NUM_THREADS", str(CPU_THREADS))

from pathlib import Path

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama


# Official LFM2.5 target model.
# Q4_0 is generally the fastest small CPU option.
MODEL_REPO = os.getenv(
    "MODEL_REPO",
    "LiquidAI/LFM2.5-2.6B-GGUF",
)

MODEL_FILE = os.getenv(
    "MODEL_FILE",
    "LFM2.5-2.6B-Q4_0.gguf",
)

# Keep this moderate on shared CPU Spaces.
# 2048 is faster than 4096 and is enough for many API requests.
N_CTX = int(os.getenv("N_CTX", "2048"))

# llama.cpp can use all visible CPUs, but shared Spaces may perform
# better with a slightly lower value. Override with CPU_THREADS.
N_THREADS = int(os.getenv("CPU_THREADS", str(CPU_COUNT)))

MODEL_PATH = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    cache_dir="/tmp/huggingface-cache",
)

print(f"Loading model: {MODEL_PATH}")
print(f"CPU threads: {N_THREADS}")
print("GPU layers: 0")
print("Backend: CPU-only")


llm = Llama(
    model_path=MODEL_PATH,

    # Absolute CPU-only settings.
    n_gpu_layers=0,
    split_mode=0,
    main_gpu=0,

    # CPU parallelism.
    n_threads=N_THREADS,
    n_threads_batch=N_THREADS,

    # Prompt-processing batch size.
    # Lower this to 256 if memory is limited.
    n_batch=512,

    # Context size.
    n_ctx=N_CTX,

    # Memory/performance settings.
    use_mmap=True,
    use_mlock=False,

    # Do not use GPU-oriented KV-cache settings.
    offload_kqv=False,
    flash_attn=False,

    # Prevent noisy native logs after startup.
    verbose=False,
)


SYSTEM_PROMPT = os.getenv(
    "SYSTEM_PROMPT",
    "You are a helpful, concise assistant.",
)


def make_prompt(user_prompt: str) -> str:
    """
    LFM2.5 understands the chat-style format stored in the GGUF metadata.
    llama-cpp-python's create_chat_completion applies the model template.
    """
    return user_prompt.strip()


def generate(user_prompt: str) -> str:
    if not user_prompt or not user_prompt.strip():
        return "Please enter a message."

    response = llm.create_chat_completion(
        messages=[
            {
                "role": "system",
                "content": SYSTEM_PROMPT,
            },
            {
                "role": "user",
                "content": make_prompt(user_prompt),
            },
        ],

        # Generation settings.
        max_tokens=int(os.getenv("MAX_TOKENS", "512")),
        temperature=float(os.getenv("TEMPERATURE", "0.2")),
        top_p=float(os.getenv("TOP_P", "0.9")),
        top_k=int(os.getenv("TOP_K", "40")),
        repeat_penalty=float(os.getenv("REPEAT_PENALTY", "1.05")),

        # Avoid unnecessary response metadata.
        stream=False,
    )

    return response["choices"][0]["message"]["content"].strip()


demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(
        label="Prompt",
        placeholder="Ask something...",
        lines=5,
    ),
    outputs=gr.Textbox(
        label="Response",
        lines=12,
    ),
    title="LFM2.5 2.6B CPU API",
    description="LFM2.5 running through a prebuilt CPU llama.cpp wheel.",
    api_name="chat",
)


if __name__ == "__main__":
    demo.queue(
        max_size=16,
        default_concurrency_limit=1,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_api=True,
    )