engine/cpu_worker.py

"""Internal JSON-lines worker. Never import spaces here; GPU visibility is disabled."""
from __future__ import annotations
import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["NVIDIA_VISIBLE_DEVICES"] = ""
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

from contextlib import redirect_stdout
from dataclasses import asdict
import json
from pathlib import Path
import sys
import numpy as np

from .config import SETTINGS

PIPE = sys.stdout
EMBEDDER = None
WHISPER = None
CPU_LLM = None


def send(message: dict) -> None:
    PIPE.write(json.dumps(message, ensure_ascii=False) + "\n")
    PIPE.flush()


def private_path(value: str) -> Path:
    p = Path(value).resolve()
    if not p.is_relative_to(SETTINGS.root) or not p.is_file() or p.is_symlink():
        raise ValueError("The temporary input file is unavailable.")
    return p


def validate_waveform(audio: np.ndarray, max_seconds: int = SETTINGS.max_audio_seconds) -> np.ndarray:
    audio = np.asarray(audio, dtype="float32")
    if audio.ndim != 1 or len(audio) < 1600:
        raise ValueError("The recording is empty or too short.")
    if len(audio) > max_seconds * 16000:
        raise ValueError(f"Record no more than {max_seconds} seconds per answer.")
    if not np.isfinite(audio).all():
        raise ValueError("Invalid audio samples.")
    if float(np.sqrt(np.mean(audio * audio))) < 0.00001:
        raise ValueError("No audible speech was detected. Your existing answer has been kept.")
    return audio


def decode_audio(path: Path) -> np.ndarray:
    """Bounded PyAV decode; do not allocate an unbounded waveform before duration validation."""
    import av
    if path.stat().st_size > SETTINGS.max_audio_bytes:
        raise ValueError("The recording is too large.")
    result = []
    samples = 0
    limit = SETTINGS.max_audio_seconds * 16000
    with av.open(str(path)) as container:
        if not container.streams.audio:
            raise ValueError("No audio track was found.")
        stream = container.streams.audio[0]
        resampler = av.audio.resampler.AudioResampler(format="fltp", layout="mono", rate=16000)
        for frame in container.decode(stream):
            for output in resampler.resample(frame):
                x = output.to_ndarray().reshape(-1).astype("float32")
                samples += len(x)
                if samples > limit:
                    raise ValueError(f"Record no more than {SETTINGS.max_audio_seconds} seconds per answer.")
                result.append(x)
        for output in resampler.resample(None):
            x = output.to_ndarray().reshape(-1).astype("float32")
            samples += len(x)
            if samples > limit:
                raise ValueError("The recording exceeds the duration limit.")
            result.append(x)
    if samples < 1600:
        raise ValueError("The recording is empty or too short.")
    return validate_waveform(np.concatenate(result))


def process(kind: str, operation: str, data: dict, request_id: str) -> dict:
    global EMBEDDER, WHISPER, CPU_LLM
    def stage(text: str) -> None:
        send({"id": request_id, "stage": text})

    if kind == "llm" and operation == "infer":
        from .rpc import validate_payload
        validated = validate_payload(data["task"], json.dumps(data["payload"]))
        if CPU_LLM is None:
            from .llm import LocalCPUBackend
            stage("Loading the local CPU language model; first use may download weights...")
            CPU_LLM = LocalCPUBackend()
        stage("Running language-model inference on CPU; this can take several minutes...")
        return json.loads(CPU_LLM._dispatch(data["task"], json.dumps(validated)))

    if kind == "documents" and operation == "parse":
        from .ingestion import parse_documents
        inputs = [(private_path(p), name, role) for p, name, role in data["files"]]
        corpus = parse_documents(inputs, data.get("warnings", []))
        return asdict(corpus)

    if kind == "documents" and operation == "embed":
        texts = data.get("texts", [])
        if not isinstance(texts, list) or not 1 <= len(texts) <= 32 or any(not isinstance(t, str) or len(t) > 15000 for t in texts):
            raise ValueError("Invalid embedding batch.")
        if EMBEDDER is None:
            stage("Loading the embedding model on CPU (first use only)...")
            from sentence_transformers import SentenceTransformer
            EMBEDDER = SentenceTransformer(SETTINGS.embedding_model, device="cpu", trust_remote_code=False, token=False)
            EMBEDDER.max_seq_length = min(getattr(EMBEDDER, "max_seq_length", 512), 512)
            print("[CPU worker] embedding model ready", file=sys.stderr, flush=True)
        values = EMBEDDER.encode(texts, batch_size=16, normalize_embeddings=True,
                                 convert_to_numpy=True, show_progress_bar=False)
        print(f"[CPU worker] embedded {len(texts)} passages", file=sys.stderr, flush=True)
        return {"vectors": values.astype("float32").tolist()}

    if kind == "speech" and operation == "transcribe":
        path = private_path(data["path"])
        stage("Decoding the recording and checking its duration...")
        audio = decode_audio(path)
        if WHISPER is None:
            stage("Loading the speech model on CPU (first use only)...")
            from faster_whisper import WhisperModel
            WHISPER = WhisperModel(SETTINGS.speech_model, device="cpu", compute_type="int8", cpu_threads=2, num_workers=1)
            print("[CPU worker] speech model ready", file=sys.stderr, flush=True)
        language = data.get("language", "en")
        allowed = {"en", "it", "fr", "de", "es", "pt", "nl", "ar", "zh", "ja", "auto"}
        if language not in allowed:
            raise ValueError("Unsupported language selection.")
        # Hint contains ONLY visible topic/optional learner vocabulary. Never
        # seed ASR with a model answer or the hidden omission rubric.
        prompt = str(data.get("vocabulary", ""))[:500].strip()
        segments, info = WHISPER.transcribe(audio, language=None if language == "auto" else language,
                                          task="transcribe", beam_size=3, vad_filter=True,
                                          condition_on_previous_text=False,
                                          initial_prompt=prompt or None)
        transcript = []
        duration = len(audio) / 16000
        for segment in segments:
            transcript.append(segment.text.strip())
            stage(f"Transcribing: {min(duration, segment.end):.0f} / {duration:.0f} seconds processed...")
        text = " ".join(t for t in transcript if t).strip()
        if not text:
            raise ValueError("No speech was detected. Your existing answer has been kept.")
        return {"text": text, "language": info.language, "seconds": round(duration, 1)}
    raise ValueError("Unsupported CPU operation.")


def main() -> None:
    kind = sys.argv[1]
    for line in sys.stdin:
        request_id = ""
        try:
            if len(line) > 2_000_000:
                raise ValueError("Request size limit exceeded.")
            request = json.loads(line)
            request_id = request["id"]
            # Third-party messages must not pollute the machine-readable pipe.
            with redirect_stdout(sys.stderr):
                result = process(kind, request["op"], request["data"], request_id)
            send({"id": request_id, "result": result})
        except Exception as exc:
            # Never write document text, answer text, or transcripts to public logs.
            from .ingestion import IngestionError
            message = str(exc) if isinstance(exc, (ValueError, IngestionError)) else f"{type(exc).__name__} in CPU worker. Check dependencies/model availability and retry."
            send({"id": request_id, "error": message[:500]})
            print(f"[CPU worker] operation failed: {type(exc).__name__}", file=sys.stderr, flush=True)


if __name__ == "__main__":
    main()