File size: 5,130 Bytes
9a1014e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
154
155
156
157
158
159
160
161
"""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"))