"""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()