hf_qwen_client.py

"""Inferencia local con el GGUF Qwen3-4B-Instruct mediante llama.cpp."""

from __future__ import annotations

import os
import threading
import time
from pathlib import Path
from typing import Iterator

from dotenv import load_dotenv
from llama_cpp import Llama


load_dotenv()

DEFAULT_MODEL = "AlbertiTechnology/qwen3-4b-instruct-gguf"
DEFAULT_MODEL_PATH = (
    Path(__file__).resolve().parent
    / "MODELS"
    / "qwen3-4b-instruct-gguf"
    / "model.gguf"
)
MODEL_ALIASES = {
    "Qwen/Qwen3-4B": DEFAULT_MODEL,
    "Qwen/Qwen3-4B-Instruct-2507": DEFAULT_MODEL,
    "Qwen/Qwen2.5-7B-Instruct": DEFAULT_MODEL,
}

_load_lock = threading.Lock()
_generation_lock = threading.Lock()
_model: Llama | None = None


def local_model_path() -> Path:
    return Path(os.getenv("LOCAL_MODEL_PATH", str(DEFAULT_MODEL_PATH))).resolve()


def _load_local_model() -> Llama:
    global _model
    if _model is not None:
        return _model

    with _load_lock:
        if _model is not None:
            return _model

        model_path = local_model_path()
        if not model_path.is_file():
            raise RuntimeError(f"Local GGUF model was not found at {model_path}")

        print(f"Loading local GGUF model from {model_path}...", flush=True)
        _model = Llama(
            model_path=str(model_path),
            n_ctx=int(os.getenv("LOCAL_MODEL_CONTEXT_SIZE", "2048")),
            n_threads=int(os.getenv("LOCAL_MODEL_THREADS", "2")),
            n_threads_batch=int(os.getenv("LOCAL_MODEL_BATCH_THREADS", "2")),
            n_batch=int(os.getenv("LOCAL_MODEL_BATCH_SIZE", "256")),
            n_gpu_layers=int(os.getenv("LOCAL_MODEL_GPU_LAYERS", "0")),
            verbose=os.getenv("LLAMA_CPP_VERBOSE", "false").lower() == "true",
        )
        print("Local GGUF model loaded.", flush=True)
        return _model


def generate_response(
    text: str,
    model: str = DEFAULT_MODEL,
    max_new_tokens: int | None = None,
) -> str:
    """Genera una respuesta usando exclusivamente el GGUF local."""
    return "".join(generate_response_stream(text, model, max_new_tokens)).strip()


def generate_response_stream(
    text: str,
    model: str = DEFAULT_MODEL,
    max_new_tokens: int | None = None,
) -> Iterator[str]:
    """Genera la respuesta local y entrega cada fragmento apenas esta disponible."""
    del model
    local_model = _load_local_model()
    token_limit = max_new_tokens or int(
        os.getenv("LOCAL_MODEL_MAX_NEW_TOKENS", "256")
    )
    context_size = int(os.getenv("LOCAL_MODEL_CONTEXT_SIZE", "2048"))
    prompt_budget = max(256, context_size - token_limit - 192)
    prompt_tokens = local_model.tokenize(
        text.encode("utf-8"),
        add_bos=False,
        special=True,
    )
    if len(prompt_tokens) > prompt_budget:
        marker = "\n\n[Contexto intermedio recortado por límite de tokens]\n\n"
        marker_tokens = local_model.tokenize(
            marker.encode("utf-8"),
            add_bos=False,
            special=True,
        )
        available = max(1, prompt_budget - len(marker_tokens))
        start_count = int(available * 0.6)
        end_count = available - start_count
        kept_tokens = (
            prompt_tokens[:start_count]
            + marker_tokens
            + prompt_tokens[-end_count:]
        )
        text = local_model.detokenize(kept_tokens).decode(
            "utf-8", errors="ignore"
        )
        print(
            "Local prompt truncated | "
            f"original_tokens={len(prompt_tokens)} | kept_tokens={len(kept_tokens)} "
            f"| budget={prompt_budget}",
            flush=True,
        )
    messages = [
        {
            "role": "system",
            "content": (
                "Responde en español de forma directa y concisa. Usa el contexto "
                "recuperado de Chroma como fuente principal. Si el contexto no "
                "alcanza, dilo claramente. No muestres razonamiento interno."
            ),
        },
        {"role": "user", "content": f"/no_think\n{text}"},
    ]

    with _generation_lock:
        started_at = time.monotonic()
        first_chunk = True
        print(
            f"Local generation prompt submitted | max_tokens={token_limit}",
            flush=True,
        )
        response = local_model.create_chat_completion(
            messages=messages,
            max_tokens=token_limit,
            temperature=0.7,
            top_p=0.8,
            top_k=20,
            repeat_penalty=1.05,
            stream=True,
        )
        for chunk in response:
            delta = chunk["choices"][0].get("delta", {})
            content = delta.get("content") or delta.get("reasoning_content")
            if content:
                if first_chunk:
                    print(
                        "Local model produced a chunk | "
                        f"elapsed={time.monotonic() - started_at:.1f}s",
                        flush=True,
                    )
                    first_chunk = False
                yield str(content)


if __name__ == "__main__":
    print(generate_response("Responde solamente: Qwen GGUF local OK"))