John6666/llama_cpp_python_test1

1import os2import threading3from typing import Any, Dict, Iterable, List, Union4 5import gradio as gr6from huggingface_hub import hf_hub_download7 8from llama_cpp import Llama9 10# -----------------------------11# Model (HF GGUF)12# -----------------------------13MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "Qwen/Qwen2.5-0.5B-Instruct-GGUF")14MODEL_FILENAME = os.getenv("MODEL_FILENAME", "qwen2.5-0.5b-instruct-q4_k_m.gguf")15 16SYSTEM_PROMPT = os.getenv(17    "SYSTEM_PROMPT",18    "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",19)20 21# Keep modest on free CPU (KV cache grows with context).22N_CTX = int(os.getenv("N_CTX", "4096"))23 24# Generation defaults25TEMPERATURE = float(os.getenv("TEMPERATURE", "0.7"))26TOP_P = float(os.getenv("TOP_P", "0.9"))27MAX_TOKENS = int(os.getenv("MAX_TOKENS", "512"))28 29# -----------------------------30# Lazy singleton model loader31# -----------------------------32_llm: Llama | None = None33_llm_lock = threading.Lock()34 35 36def _load_llm() -> Llama:37    global _llm38    if _llm is not None:39        return _llm40 41    with _llm_lock:42        if _llm is not None:43            return _llm44 45        model_path = hf_hub_download(repo_id=MODEL_REPO_ID, filename=MODEL_FILENAME)46 47        # Qwen instruct GGUFs commonly use ChatML-style formatting.48        _llm = Llama(49            model_path=model_path,50            n_ctx=N_CTX,51            n_threads=os.cpu_count() or 4,52            n_gpu_layers=0,53            chat_format="chatml",54            verbose=False,55        )56        return _llm57 58 59# -----------------------------60# Gradio message normalization61# -----------------------------62Content = Union[str, List[Any], Dict[str, Any]]63 64 65def _content_to_text(content: Content) -> str:66    if isinstance(content, str):67        return content68    if isinstance(content, list):69        parts: List[str] = []70        for item in content:71            if isinstance(item, str):72                parts.append(item)73            elif isinstance(item, dict) and item.get("type") == "text":74                parts.append(str(item.get("text", "")))75        return "".join(parts).strip()76    if isinstance(content, dict):77        for k in ("text", "content"):78            v = content.get(k)79            if isinstance(v, str):80                return v81    return str(content)82 83 84def _history_to_messages(history: Any) -> List[Dict[str, str]]:85    if not history:86        return []87 88    msgs: List[Dict[str, str]] = []89 90    # Old format: list[(user, assistant), ...]91    if isinstance(history, list) and history and isinstance(history[0], (tuple, list)) and len(history[0]) == 2:92        for user, assistant in history:93            if user:94                msgs.append({"role": "user", "content": str(user)})95            if assistant:96                msgs.append({"role": "assistant", "content": str(assistant)})97        return msgs98 99    # Newer format: list[{"role": "...", "content": ...}, ...]100    if isinstance(history, list) and history and isinstance(history[0], dict):101        for m in history:102            role = m.get("role")103            if role not in ("user", "assistant", "system"):104                continue105            text = _content_to_text(m.get("content", ""))106            if text:107                msgs.append({"role": role, "content": text})108        return msgs109 110    return []111 112 113def _stream_chat(llm: Llama, messages: List[Dict[str, str]]) -> Iterable[str]:114    # llama-cpp-python yields OpenAI-like streaming chunks.115    stream = llm.create_chat_completion(116        messages=messages,117        temperature=TEMPERATURE,118        top_p=TOP_P,119        max_tokens=MAX_TOKENS,120        stream=True,121    )122 123    partial = ""124    for chunk in stream:125        token = ""126        try:127            choice = chunk["choices"][0]128            delta = choice.get("delta") or {}129            token = delta.get("content") or ""130        except Exception:131            token = ""132        if token:133            partial += token134            yield partial135 136 137def respond(message: str, history: Any):138    llm = _load_llm()139 140    msgs: List[Dict[str, str]] = [{"role": "system", "content": SYSTEM_PROMPT}]141    prior = _history_to_messages(history)142 143    # Simple history trim144    if len(prior) > 20:145        prior = prior[-20:]146 147    msgs.extend(prior)148    msgs.append({"role": "user", "content": message})149 150    for partial in _stream_chat(llm, msgs):151        yield partial152 153 154demo = gr.ChatInterface(155    fn=respond,156    title="GGUF Chatbot (llama-cpp-python)",157)158 159if __name__ == "__main__":160    demo.launch()161