"""EuropaLex Inference Engine — Local model backends via llama-cli, llama-cpp-python, and Python packages.""" from __future__ import annotations import logging import os import re from dataclasses import dataclass from pathlib import Path from typing import ClassVar import numpy as np import torch from core.types import CEFRLevel, EngineConfig, TextResult, ValidationError logger = logging.getLogger(__name__) # Global pre-loaded engines for zeroGPU module-level loading _engines: dict[str, Any] = {} _preload_success = False class _LazyEngineWrapper: """Lazy-loading fallback wrapper for engines that failed to pre-load. In HF Spaces, torchaudio/torch mismatches may prevent TTS/Image engines from loading. This wrapper will attempt on-demand but fail gracefully if the underlying import still fails. """ def __init__(self, name: str, factory_func): self.name = name self.factory = factory_func self._engine = None def _get_engine(self): if self._engine is None: try: engine = self.factory() # Load the model into CUDA (emulation mode -> real GPU at runtime) if hasattr(engine, '_load_model'): engine._load_model() elif hasattr(engine, '_load_pipeline'): engine._load_pipeline() self._engine = engine logger.info("Loaded fallback %s engine", self.name) except ImportError as e: # For HF Space environment issues, provide clearer error if 'torch_library_impl' in str(e) or 'torchaudio' in str(e): raise RuntimeError( f"Failed to load {self.name} engine due to torch/torchaudio version mismatch. " f"This is common in Hugging Face Spaces with custom CUDA builds.\n" f"Please report this to the maintainer or try using the app without audio/image generation." ) from e raise RuntimeError(f"Failed to load {self.name} engine: {e}") from e return self._engine def __getattr__(self, item): return getattr(self._get_engine(), item) def _preload_all_models() -> None: """Pre-load all inference models at module level. Falls back to lazy-loading if any fail.""" global _preload_success if _preload_success: return from core.types import EngineConfig try: config = EngineConfig.from_settings_yaml() except Exception as e: logger.error("Failed to load settings: %s", e) # Create fallback wrappers with minimal info _preload_success = False # will retry on first use return # HF Spaces often have torchaudio/torch version mismatches. # For TTS/Image engines, skip preloading and defer to lazy loading. # Text engines (MiniCPM/LlamaCpp) don't depend on torchaudio. is_hf_space = Path('/home/space/app').exists() or 'HF_SPACE' in os.environ engines_to_load = [ ('english', lambda: MiniCPMTextEngine(config.minicpm_model_path, config.device)), ('translation', lambda: LlamaCppTextEngine(config.llm_model_path, config.device, config.target_language)), ] if not is_hf_space: engines_to_load.extend([ ('tts', lambda: TTSEngine(device=config.device)), ('image', lambda: ImageGenEngine(device=config.device)), ]) for name, factory in engines_to_load: try: engine = factory() # Force load the model into CUDA if hasattr(engine, '_load_model'): engine._load_model() elif hasattr(engine, '_load_pipeline'): engine._load_pipeline() else: logger.warning("Engine %s has no load method", name) _engines[name] = engine logger.info("Pre-loaded %s engine", name) except Exception as e: logger.error("Failed to pre-load %s engine: %s", name, e) # Create lazy fallback wrapper for this engine _engines[name] = _LazyEngineWrapper(name, factory) # Trigger actual GPU allocation (forces emulation mode -> real GPU) import torch if torch.cuda.is_available(): try: torch.tensor([0], device='cuda') except Exception as e: logger.error("Failed to trigger CUDA allocation: %s", e) _preload_success = True logger.info("Pre-load complete. Engines: %s", list(_engines.keys())) @dataclass class _EngineState: """Tracks which GPU engine is currently loaded.""" translation_engine: LlamaCppTextEngine | None = None tts_engine: TTSEngine | None = None image_engine: ImageGenEngine | None = None class MiniCPMTextEngine: """Generates English text using MiniCPM5-1B Q8_0 via llama-cpp-python. Lazy-loads the model on first call, unloads after completion to free memory. Uses create_chat_completion to format prompts with MiniCPM's required message tokens. Only one instance can be active at a time (enforced by EnginePool). Best for Phase 1 English text generation — ~1.1 GB RAM, no subprocess overhead. """ def __init__(self, model_path: str, device: str = "cuda"): """Initialize the text engine. Args: model_path: Absolute path to the MiniCPM5-1B Q8_0 GGUF file. device: Device hint ('cuda', 'mps', or 'cpu'). """ self.model_path = Path(model_path) if not self.model_path.exists(): raise FileNotFoundError( f"MiniCPM5-1B model not found at: {self.model_path}\n" f"Run: python models/download_models.py minicpm" ) self.device = device self._llm = None self._loaded = False def _load_model(self) -> None: """Lazy-load the GGUF model via llama-cpp-python.""" if self._loaded: return try: from llama_cpp import Llama except ImportError: raise ImportError( "llama-cpp-python package not installed. " "Run: pip install llama-cpp-python" ) n_gpu = 99 if self.device == "cuda" else 0 self._llm = Llama( model_path=str(self.model_path), n_gpu_layers=n_gpu, n_ctx=4096, ) self._loaded = True logger.info("MiniCPMTextEngine loaded %s on %s", self.model_path.name, self.device) def generate( self, texts: list[str], scenario: str, cefr_level: CEFRLevel, batch_size: int | None = None, topic_description: str = "", ) -> TextResult: """Generate English sentences using the loaded GGUF model. Delegates to :func:`core.text_gen.generate_sentences` for LLM calling, retry loop, and extraction. Wraps result in ``TextResult``. Args: texts: Empty list (generation mode). Non-empty would be translation mode. scenario: Scenario/topic description for text generation. cefr_level: CEFR proficiency level (linguistic guidance only). batch_size: Number of sentences to generate. topic_description: Free-form description of topics/themes. Returns: TextResult with exactly one sentence per requested batch size. Raises: ValidationError: If generation fails after max attempts. """ self._load_model() if batch_size is None: raise ValueError("batch_size is required for text generation") from core.text_gen import generate_sentences sentences = generate_sentences(scenario, cefr_level, batch_size, self._llm, topic_description) return TextResult(generated_texts=sentences) class LlamaCppTextEngine: """Generates text using llama-cpp-python (GGUF models, lazy-load + unload). Uses the llama-cpp-python Python bindings instead of spawning subprocesses. Lazy-loads the model on first call, unloads after completion to free VRAM. Only one instance can be active at a time (enforced by EnginePool). Best for smaller GGUF models (e.g. tiny-aya-water Q4_K_M ~2 GB) where keeping the model in Python memory is efficient and avoids subprocess overhead. """ def __init__( self, model_path: str, device: str = "cuda", target_language: str = "Latvian", ): """Initialize the translation engine. Args: model_path: Absolute path to the GGUF model file. device: Device hint (informational; n_gpu_layers=99 used for CUDA). target_language: Target language for translations (e.g. "Latvian"). """ self.model_path = Path(model_path) if not self.model_path.exists(): raise FileNotFoundError(f"Model not found: {self.model_path}") self.device = device self.target_language = target_language self._llm = None self._loaded = False def _load_model(self) -> None: """Lazy-load the GGUF model via llama-cpp-python.""" if self._loaded: return try: from llama_cpp import Llama except ImportError: raise ImportError( "llama-cpp-python package not installed. " "Run: pip install llama-cpp-python" ) n_gpu = 99 if self.device == "cuda" else 0 self._llm = Llama( model_path=str(self.model_path), n_gpu_layers=n_gpu, n_ctx=4096, ) self._loaded = True logger.info("LlamaCppTextEngine loaded %s on %s", self.model_path.name, self.device) def _translate_single( self, text: str, cefr_level: CEFRLevel, topic_description: str = "", target_language: str | None = None, ) -> str: """Translate a single English sentence with retry loop. Uses ``create_chat_completion`` so the model's chat template (with special tokens like ``<|USER_TOKEN|>``) is applied correctly. Raw prompt strings bypass the chat template and produce poor output. Wraps the LLM call in a retry loop (max 3 attempts). Returns the translated string or falls back to the original English text on failure. Args: text: Single English sentence to translate. cefr_level: CEFR proficiency level (linguistic guidance only). topic_description: Free-form description of topics/themes (for context). target_language: Override target language for this call only. Defaults to ``self.target_language`` from config. Returns: Translated string, or the original English text as fallback. """ self._load_model() effective_lang = target_language or self.target_language base_messages = [ { "role": "system", "content": ( f"You are a professional translator. Translate English sentences into " f"{effective_lang} at the specified CEFR level. Output ONLY the translation, " f"one line. No explanations, no notes, no source text repetition." ), } ] last_messages: list = [] for attempt in range(1, 4): messages = list(base_messages) if attempt > 1 and last_messages: # Append failed output + retry instruction in conversation context messages.extend(last_messages) messages.append({ "role": "user", "content": self._build_single_translation_prompt( text, cefr_level, topic_description, effective_lang, ), }) try: output = self._llm.create_chat_completion( messages=messages, max_tokens=128, temperature=0.3, ) except Exception as e: # CUDA decode failures (e.g. ggml-cuda launch_mul_mat_q on SWA models) # are often fixed by reloading the model — the fresh context gets # proper CUDA graph warmup. logger.warning( "LlamaCppTextEngine: decode failed on attempt %d for '%s': %s — reloading model and retrying", attempt, text[:30], e, ) self.unload() # free GPU memory from stale CUDA state self._load_model() # fresh load gets clean warmup continue raw_text = output.get("choices", [{}])[0].get("message", {}).get("content", "") last_messages = [ {"role": "assistant", "content": raw_text}, ] line = raw_text.strip() # Validate: must be a single non-empty line, not repetitive garbage if line and self._is_valid_translation(line): logger.info( "LlamaCppTextEngine: translated '%s' on attempt %d -> '%s'", text[:30], attempt, line[:40], ) return line # Invalid output — retry with stricter prompt in context if attempt < 3: last_messages.append({ "role": "user", "content": ( f"That output was invalid. Translate ONLY this sentence into {effective_lang}:\n" f"{text}\n\n" f"Output ONE line only — the translation. Nothing else." ), }) logger.warning( "LlamaCppTextEngine attempt %d: invalid output for '%s' — retrying", attempt, text[:30], ) else: logger.warning( "LlamaCppTextEngine: exhausted retries for '%s'. Falling back to English.", text[:30], ) # Exhausted retries — fall back to original English text logger.info("LlamaCppTextEngine: fallback to English for '%s'", text[:30]) return text def _is_valid_translation(self, line: str) -> bool: """Check if a translation output is valid (single line, not repetitive garbage). Args: line: The raw model output to validate. Returns: True if the output looks like a reasonable translation. """ if not line or len(line) < 2: return False # Reject if it contains multiple lines (model generated too much) if "\n" in line: return False # Reject if it's just the English text back (no translation happened) lower = line.lower() if any(word in lower for word in ["translate", "translation", "english"]): return False # Reject very short outputs that are likely noise words = line.split() if len(words) < 1: return False return True def generate( self, texts: list[str], scenario: str, cefr_level: CEFRLevel, batch_size: int | None = None, topic_description: str = "", ) -> TextResult: """Generate translations using the loaded GGUF model. Translates each sentence individually for better quality with small models. Falls back to English text on failure (after max retries). Args: texts: English sentences to translate. scenario: Scenario/topic description (contextual). cefr_level: CEFR proficiency level (linguistic guidance only). batch_size: Number of translations expected. topic_description: Free-form description of topics/themes (contextual). Returns: TextResult with one translation per input text. Raises: ValidationError: If generation fails after max attempts and no lines produced. """ self._load_model() if batch_size is None: raise ValueError("batch_size is required for translation") translations = [] for text in texts: translated = self._translate_single(text, cefr_level, topic_description) translations.append(translated) return TextResult(generated_texts=translations) def _build_single_translation_prompt( self, text: str, cefr_level: CEFRLevel, topic_description: str = "", target_language: str | None = None, ) -> str: """Build prompt for translating a single sentence. Optimized for small models (tiny-aya-water ~3.3B params). Produces natural, idiomatic output by emphasizing how native speakers actually phrase things — not literal word-for-word translation. Uses CEFR linguistic guidance only — no hardcoded topics. Args: text: English sentence to translate. cefr_level: CEFR proficiency level. topic_description: Free-form context for the translation. target_language: Language to translate into. Defaults to ``self.target_language``. """ target_lang = target_language or self.target_language cefr_desc = cefr_level.description() topic_hint = f" Context: {topic_description}." if topic_description else "" return ( f"Translate the following English sentence into {target_lang}.{topic_hint}\n" f"CEFR linguistic guidance: {cefr_desc}.\n\n" f"CRITICAL — NATURAL LANGUAGE RULES:\n" f"1. Produce how a native speaker at this CEFR level would naturally express this idea in {target_lang}.\n" f"2. Do NOT translate word-for-word. Capture the meaning and rephrase it naturally in {target_lang}.\n" f"3. Use common idiomatic expressions, colloquial phrasing, and everyday vocabulary appropriate for the level.\n" f"4. Follow the grammar patterns typical of {target_lang} — not English sentence structure.\n" f"5. If the English uses an awkward or literal construction, render it as a native speaker would say it.\n\n" f"Rules:\n" f"1. Output ONLY the translated sentence — one line, nothing else.\n" f"2. Do NOT include explanations, notes, labels, or quotation marks.\n" f"3. Do NOT repeat the English text in your output.\n" f"4. Match the CEFR linguistic complexity for the target level.\n\n" f"English: {text}\n" f"{target_lang}:" ) # TTSEngine has been extracted to core.audio_gen from core.audio_gen import TTSEngine # noqa: F401 # ImageGenEngine has been extracted to core.image_gen from core.image_gen import ImageGenEngine # noqa: F401 class EnginePool: """Singleton providing access to pre-loaded inference engines. Maintains backward compatibility with the old singleton factory pattern. Engines are pre-loaded at module import time via _preload_all_models(). """ _instance = None def __new__(cls): if cls._instance is None: raise RuntimeError( "EnginePool must be accessed via EnginePool.get(config), not directly." ) return cls._instance @classmethod def get(cls, config: EngineConfig) -> "EnginePool": """Get or create the singleton EnginePool instance. Args: config: EngineConfig (not used for instantiation but kept for compatibility). Returns: The singleton EnginePool instance. """ if cls._instance is None: instance = super().__new__(cls) instance._config = config cls._instance = instance return cls._instance def get_english_engine(self) -> MiniCPMTextEngine: """Get the pre-loaded English text generation engine.""" return _engines['english'] def get_translation_engine(self) -> LlamaCppTextEngine: """Get the pre-loaded translation engine.""" return _engines['translation'] def get_tts_engine(self) -> TTSEngine: """Get the pre-loaded TTS engine.""" return _engines['tts'] def get_image_engine(self) -> ImageGenEngine: """Get the pre-loaded image generation engine.""" return _engines['image'] # Pre-load all models at module import time _preload_all_models() # Check if TTS/Image are available (for UI conditional rendering) def _check_media_availability() -> tuple[bool, bool]: """Check if TTS and Image generation engines are functional. Returns: Tuple of (tts_available, image_available) booleans. """ import os from pathlib import Path is_hf_space = Path('/home/space/app').exists() or 'HF_SPACE' in os.environ if is_hf_space: # HF Spaces often have torchaudio/torch mismatches return False, False try: import torchaudio import torch # Basic sanity check - can we import the modules? return True, True except Exception: return False, False _TTS_AVAILABLE, _IMAGE_AVAILABLE = _check_media_availability()