ViniciusKhan/analytical_force

1"""Cliente Hugging Face Transformers (modelo público, opcional).2 3Carrega um modelo público de geração de texto via ``transformers`` e produz4a interpretação do relatório localmente, sem chamar nenhuma API comercial.5 6Cuidados:7- O modelo NÃO é fixado no código: vem de ``HF_MODEL_REPO_ID``.8- ``transformers``/``torch`` são dependências opcionais (import tardio).9- Erros de memória, modelo ausente ou dependência faltando levantam10  :class:`TransformersError` para acionar o fallback para template.11 12Princípio: este cliente NUNCA calcula indicadores. Apenas interpreta o prompt.13"""14 15from __future__ import annotations16 17from typing import Any18 19from ..utils.logger import get_logger20 21logger = get_logger("models.transformers")22 23# Cache de pipelines por (repo_id, device). Evita recarregar o modelo a cada24# chamada — o que era o principal custo quando há várias gerações por execução.25# Compartilhado por todas as instâncias no mesmo processo.26_PIPELINE_CACHE: dict[tuple[str, str], Any] = {}27 28 29class TransformersError(RuntimeError):30    """Erro controlado ao usar Hugging Face Transformers."""31 32 33class TransformersClient:34    """Encapsula um pipeline de ``text-generation`` da Hugging Face."""35 36    def __init__(37        self,38        repo_id: str,39        device: str = "cpu",40        max_new_tokens: int = 512,41    ) -> None:42        """Inicializa o cliente (sem carregar o modelo ainda).43 44        Args:45            repo_id: Identificador do modelo público (``HF_MODEL_REPO_ID``).46            device: ``cpu`` ou ``cuda`` (``HF_DEVICE``).47            max_new_tokens: Limite de tokens gerados na resposta.48 49        Raises:50            TransformersError: Se ``repo_id`` não for informado.51        """52        if not repo_id:53            raise TransformersError(54                "HF_MODEL_REPO_ID não configurado. Defina um modelo público válido."55            )56        self._repo_id = repo_id57        self._device = device58        self._max_new_tokens = max_new_tokens59        self._pipeline = None  # carregamento preguiçoso (lazy)60 61    def _carregar_pipeline(self):62        """Carrega (uma única vez) o pipeline de geração de texto.63 64        Reaproveita um pipeline em cache para o mesmo (modelo, device), evitando65        recarregar o modelo a cada chamada.66        """67        if self._pipeline is not None:68            return self._pipeline69 70        chave = (self._repo_id, self._device)71        em_cache = _PIPELINE_CACHE.get(chave)72        if em_cache is not None:73            self._pipeline = em_cache74            return self._pipeline75 76        try:77            from transformers import pipeline  # import tardio78        except ImportError as exc:  # pragma: no cover - depende do ambiente79            raise TransformersError(80                "Dependência ausente. Instale: pip install transformers torch"81            ) from exc82 83        try:84            # device=-1 força CPU; device=0 usa a primeira GPU.85            device_idx = 0 if self._device.lower().startswith("cuda") else -186            self._pipeline = pipeline(87                "text-generation",88                model=self._repo_id,89                device=device_idx,90            )91        except Exception as exc:  # modelo ausente, sem memória, etc.92            raise TransformersError(93                f"Falha ao carregar o modelo '{self._repo_id}': {type(exc).__name__}."94            ) from exc95        _PIPELINE_CACHE[chave] = self._pipeline96        logger.info("Modelo Transformers carregado (repo=%s).", self._repo_id)97        return self._pipeline98 99    def gerar(self, prompt: str, system: str | None = None) -> str:100        """Gera a interpretação textual a partir do prompt.101 102        Args:103            prompt: Prompt completo (com o JSON de métricas).104            system: Instrução de sistema opcional (prefixada ao prompt).105 106        Returns:107            Texto interpretativo gerado pelo modelo.108 109        Raises:110            TransformersError: Em qualquer falha de carga/geração.111        """112        pipe = self._carregar_pipeline()113        entrada = f"{system}\n\n{prompt}" if system else prompt114        try:115            saida = pipe(116                entrada,117                max_new_tokens=self._max_new_tokens,118                do_sample=False,119                return_full_text=False,120            )121        except Exception as exc:122            raise TransformersError(123                f"Falha na geração com Transformers: {type(exc).__name__}."124            ) from exc125 126        texto = ""127        if isinstance(saida, list) and saida:128            texto = str(saida[0].get("generated_text", "")).strip()129        if not texto:130            raise TransformersError("Transformers retornou resposta vazia.")131        logger.info("Interpretação gerada via Transformers (repo=%s).", self._repo_id)132        return texto133