Upload toarim_embedder.py with huggingface_hub

toarim_embedder.py

@@ -0,0 +1,510 @@
+"""
+=============================================================================
+ TÓARIM EMBEDDER - MOTOR DE REPRESENTACIÓN VECTORIAL INDUSTRIAL CON ESTADO
+ Arquitectura: GNN Secuencial O(L) + SVD Low-Rank + Memoria Continua Titan (Sherman-Morrison)
+ Optimizado para Cómputo CPU Exclusivo (HPC) y Cumplimiento de Ciberseguridad
+=============================================================================
+"""
+
+import os
+import json
+import logging
+from pathlib import Path
+from typing import List, Tuple, Optional, Union, Dict, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# Configuración de Logging
+logger = logging.getLogger("ToarimEmbedder")
+
+# =======================================================================
+# 1. ENRIQUECIMIENTO ESTRUCTURAL O(L) (GNN Lineal para Secuencias)
+# =======================================================================
+class DocumentGraphEnricher(nn.Module):
+    """
+    Aplica una convolución de grafos (GCN) sobre secuencias documentales
+    con complejidad lineal O(L) en memoria y cómputo.
+    
+    OPTIMIZACIÓN HPC:
+    Sustituye la matriz de adyacencia densa A_hat [L, L] de O(L^2) por
+    un operador de difusión tridiagonal local O(L) matemáticamente idéntico:
+      h_i = d_i * sum_{j in N(i)} (d_j * x_j * W)
+    Elimina por completo el riesgo de OOM DoS en secuencias extensas.
+    """
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+        self.W_gnn = nn.Linear(input_dim, output_dim, bias=False)
+        self.activation = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: Tensor de shape [batch_size, seq_len, input_dim]
+        Retorna: Tensor de shape [batch_size, seq_len, output_dim]
+        """
+        batch_size, seq_len, _ = x.shape
+        device = x.device
+        
+        # 1. Proyección lineal: [B, L, D_out]
+        x_trans = self.W_gnn(x)
+        
+        if seq_len == 1:
+            return self.activation(x_trans)
+            
+        # 2. Grados del grafo secuencial con auto-lazos:
+        # Extremos tienen grado 2 (self + 1 vecino); nodos internos tienen grado 3 (self + 2 vecinos)
+        degrees = torch.full((seq_len,), 3.0, device=device, dtype=x.dtype)
+        degrees[0] = 2.0
+        degrees[-1] = 2.0
+        d_inv_sqrt = 1.0 / torch.sqrt(degrees)  # [seq_len]
+        
+        # 3. Difusión normalizada espectralmente en O(L) memoria
+        # z_j = d_j * x_trans_j
+        z = x_trans * d_inv_sqrt.unsqueeze(0).unsqueeze(2)  # [B, L, D]
+        
+        # Suma de vecinos: s_i = z_{i-1} + z_i + z_{i+1}
+        s = z.clone()
+        s[:, :-1, :] += z[:, 1:, :]  # Conexión hacia adelante
+        s[:, 1:, :] += z[:, :-1, :]  # Conexión hacia atrás
+        
+        # h_i = d_i * s_i
+        h_graph = s * d_inv_sqrt.unsqueeze(0).unsqueeze(2)
+        return self.activation(h_graph)
+
+# =======================================================================
+# 2. COMPRESIÓN DE ESPACIO LATENTE (Low-Rank SVD Projection)
+# =======================================================================
+class LowRankLinear(nn.Module):
+    """
+    Capa lineal parametrizada mediante factorización SVD: W approx U @ V.
+    Reduce la complejidad de parámetros y memoria de O(d_in * d_out) a O(k * (d_in + d_out)).
+    """
+    def __init__(self, in_features: int, out_features: int, rank: int):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.rank = rank
+        
+        self.U = nn.Parameter(torch.randn(out_features, rank) / np.sqrt(rank))
+        self.V = nn.Parameter(torch.randn(rank, in_features) / np.sqrt(in_features))
+        self.bias = nn.Parameter(torch.zeros(out_features))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        W_approx = self.U @ self.V
+        return F.linear(x, W_approx, self.bias)
+
+# =======================================================================
+# 3. BUCLE INTERNO TITAN MEMORY (Sherman-Morrison Estable)
+# =======================================================================
+class TitanMemoryLayer(nn.Module):
+    """
+    Capa de Memoria Recurrente Asociativa basada en Nested Learning.
+    Actualización matricial O(k^2) por paso mediante el Lema de Sherman-Morrison.
+    """
+    def __init__(self, dim: int, latent_k: int):
+        super().__init__()
+        self.latent_k = latent_k
+        rank_proj = max(1, latent_k // 2)
+        
+        self.proj_q = LowRankLinear(dim, latent_k, rank=rank_proj)
+        self.proj_k = LowRankLinear(dim, latent_k, rank=rank_proj)
+        self.proj_v = LowRankLinear(dim, latent_k, rank=rank_proj)
+        
+        self.gate_eta = nn.Sequential(nn.Linear(dim, 1), nn.Sigmoid())
+        self.gate_alpha = nn.Sequential(nn.Linear(dim, 1), nn.Sigmoid())
+        self.proj_out = nn.Linear(latent_k, dim)
+
+    def forward(self, x: torch.Tensor, initial_state: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, _ = x.shape
+        device = x.device
+        
+        if initial_state is not None:
+            M = initial_state
+        else:
+            M = torch.zeros((batch_size, self.latent_k, self.latent_k), device=device, dtype=x.dtype)
+            
+        outputs = []
+        
+        for t in range(seq_len):
+            x_t = x[:, t, :]  # [B, dim]
+            
+            q_t = self.proj_q(x_t)  # [B, k]
+            k_t = self.proj_k(x_t)
+            v_t = self.proj_v(x_t)
+            
+            eta_t = self.gate_eta(x_t)      # [B, 1]
+            alpha_t = self.gate_alpha(x_t)  # [B, 1]
+            
+            # Inferencia asociativa: y_t = M * q_t
+            y_t = torch.bmm(M, q_t.unsqueeze(2)).squeeze(2)
+            
+            # Error / Sorpresa: e_t = v_t - M * k_t
+            m_k = torch.bmm(M, k_t.unsqueeze(2)).squeeze(2)
+            e_t = v_t - m_k
+            
+            # Sherman-Morrison: Denominador con regularización epsilon para estabilidad numérica
+            k_norm_sq = torch.sum(k_t ** 2, dim=1, keepdim=True)
+            denom = 1.0 + (eta_t * k_norm_sq) + 1e-7
+            update_factor = (eta_t / denom).unsqueeze(2)  # [B, 1, 1]
+            
+            outer_prod = torch.bmm(e_t.unsqueeze(2), k_t.unsqueeze(1))  # [B, k, k]
+            M = (alpha_t.unsqueeze(2) * M) + (update_factor * outer_prod)
+            
+            outputs.append(y_t)
+            
+        Y = torch.stack(outputs, dim=1)  # [B, seq_len, k]
+        return self.proj_out(Y), M
+
+# =======================================================================
+# 4. MODELO COMPLETO TÓARIM (Graph + Titan + Residual Pooling)
+# =======================================================================
+class TóarimEmbeddingModel(nn.Module):
+    """
+    Modelo de Embedding Integral Tóarim.
+    Token Embedding -> GNN Secuencial O(L) -> Titan Memory -> Residual LayerNorm -> Mean Pooling -> L2 Norm.
+    """
+    def __init__(self, vocab_size: int = 30522, embed_dim: int = 768, latent_k: int = 128):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.latent_k = latent_k
+        
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.gnn_enricher = DocumentGraphEnricher(embed_dim, embed_dim)
+        self.titan_memory = TitanMemoryLayer(embed_dim, latent_k)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, input_ids: torch.Tensor, initial_state: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = self.token_embedding(input_ids)
+        x_graph = self.gnn_enricher(x)
+        x_titan, M_final = self.titan_memory(x_graph, initial_state)
+        x_norm = self.layer_norm(x_titan + x)
+        
+        doc_embedding = torch.mean(x_norm, dim=1)
+        normalized_emb = F.normalize(doc_embedding, p=2, dim=1)
+        return normalized_emb, M_final
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs) -> "TóarimEmbeddingModel":
+        """Carga los pesos desde model.safetensors (local o Hugging Face Hub) en el módulo PyTorch."""
+        embedder = TóarimEmbedder.from_pretrained(pretrained_model_name_or_path, **kwargs)
+        return embedder.model
+
+# =======================================================================
+# 5. API DE PRODUCCIÓN TÓARIM (HPC + Ciberseguridad + LangChain/HF)
+# =======================================================================
+class TóarimEmbedder:
+    """
+    Interfaz de producción para TóarimEmbedder.
+    
+    MEDIDAS DE CIBERSEGURIDAD:
+    - Serialización segura exclusiva con Safetensors (previene vulnerabilidades de RCE por Pickle).
+    - Despliegue Air-Gapped nativo: Detección y carga local de tokenizer.json sin salida a Internet.
+    - Sanitización y aislamiento de memoria multi-inquilino mediante reset_state().
+    - Validación y truncamiento de entradas para prevenir ataques DoS por agotamiento de RAM.
+    
+    OPTIMIZACIONES HPC EN CPU:
+    - Convolución secuencial O(L) de memoria plana sin sobrecarga cuadrática.
+    - Aceleración con torch.autocast(dtype=torch.bfloat16) en procesadores compatibles (AVX-512_BF16 / AMX).
+    - Soporte integrado de cuantización dinámica INT8 (quantize()) para VNNI.
+    - Tokenización multi-hilo nativa en Rust sin sobrecarga del runtime de transformers.
+    """
+    def __init__(
+        self,
+        vocab_size: int = 30522,
+        embed_dim: int = 768,
+        latent_k: int = 128,
+        model_path: Optional[str] = None,
+        compile_model: bool = False,
+        device: Optional[str] = None,
+        tokenizer_name_or_path: str = "bert-base-uncased",
+        local_files_only: bool = False
+    ):
+        self.vocab_size = vocab_size
+        self.embed_dim = embed_dim
+        self.latent_k = latent_k
+        self.is_quantized = False
+        self._current_memory_state: Optional[torch.Tensor] = None
+
+        # Selecci��n de dispositivo y optimización de hilos en CPU
+        if device is None:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        else:
+            self.device = torch.device(device)
+            
+        if self.device.type == 'cpu':
+            # Evita sobre-suscripción y contención de hilos en servidores multi-socket NUMA
+            if hasattr(torch, "get_num_interop_threads") and torch.get_num_interop_threads() > 1:
+                try:
+                    torch.set_num_interop_threads(1)
+                except RuntimeError:
+                    pass
+
+        # -------------------------------------------------------------------
+        # Tokenizador Seguro y Air-Gapped (Carga Local Prioritaria)
+        # -------------------------------------------------------------------
+        from tokenizers import Tokenizer
+        
+        # 1. Verificar si existe tokenizer.json en la ruta indicada o en el directorio local
+        local_candidate = Path(tokenizer_name_or_path)
+        if local_candidate.is_dir() and (local_candidate / "tokenizer.json").exists():
+            tok_path = str(local_candidate / "tokenizer.json")
+            self.tokenizer = Tokenizer.from_file(tok_path)
+        elif local_candidate.is_file():
+            self.tokenizer = Tokenizer.from_file(str(local_candidate))
+        elif Path("tokenizer.json").exists():
+            # Fallback seguro offline en la raíz del proyecto
+            self.tokenizer = Tokenizer.from_file("tokenizer.json")
+        elif not local_files_only:
+            # Descarga remota autorizada solo si no es entorno air-gapped estricto
+            self.tokenizer = Tokenizer.from_pretrained(tokenizer_name_or_path)
+        else:
+            raise FileNotFoundError(
+                f"Modo Air-Gapped activado: No se localizó 'tokenizer.json' en '{tokenizer_name_or_path}'."
+            )
+            
+        self.tokenizer.no_padding()
+        self.tokenizer.no_truncation()
+
+        # Inicializar modelo
+        self.model = TóarimEmbeddingModel(vocab_size, embed_dim, latent_k).to(self.device)
+        
+        # -------------------------------------------------------------------
+        # Carga Segura de Pesos (Safetensors / Weights-Only)
+        # -------------------------------------------------------------------
+        if model_path:
+            self._load_weights_safely(model_path)
+            
+        self.model.eval()
+
+        if compile_model:
+            try:
+                self.model = torch.compile(self.model)
+            except Exception as e:
+                logger.warning(f"No fue posible aplicar torch.compile: {e}. Continuando en modo estándar.")
+
+    def _load_weights_safely(self, model_path: str):
+        """Carga pesos aplicando blindaje contra deserialización maliciosa."""
+        path = Path(model_path)
+        if not path.exists():
+            raise FileNotFoundError(f"Archivo de pesos no encontrado: {model_path}")
+            
+        if path.suffix == ".safetensors":
+            from safetensors.torch import load_file
+            state_dict = load_file(str(path), device=str(self.device))
+            self.model.load_state_dict(state_dict)
+        else:
+            # Blindaje contra RCE en archivos .pt/.bin mediante weights_only=True
+            state_dict = torch.load(str(path), map_location=self.device, weights_only=True)
+            self.model.load_state_dict(state_dict)
+
+    # ===================================================================
+    # Métodos de Ciberseguridad: Sanitización de Memoria
+    # ===================================================================
+
+    def reset_state(self):
+        """
+        Higienización y sanitización operativa del estado de memoria residual M.
+        OBLIGATORIO en entornos multi-inquilino o al cambiar entre manuales técnicos
+        para prevenir la fuga de información confidencial (Cross-Tenant Memory Leak).
+        """
+        self._current_memory_state = None
+        if self.device.type == 'cuda':
+            torch.cuda.empty_cache()
+
+    # ===================================================================
+    # Métodos de HPC: Cuantización Dinámica INT8 en CPU
+    # ===================================================================
+
+    def quantize(self):
+        """
+        Aplica cuantización dinámica INT8 sobre capas lineales.
+        Acelera la inferencia en CPUs modernas x86/ARM con instrucciones VNNI/AMX,
+        reduciendo el consumo de RAM en ~50%.
+        """
+        if self.device.type != 'cpu':
+            logger.warning("La cuantización INT8 está recomendada exclusivamente para inferencia en CPU.")
+            return self
+            
+        self.model = torch.ao.quantization.quantize_dynamic(
+            self.model,
+            {nn.Linear, LowRankLinear},
+            dtype=torch.qint8
+        )
+        self.is_quantized = True
+        return self
+
+    # ===================================================================
+    # Inferencia y Codificación de Vectores
+    # ===================================================================
+
+    @torch.no_grad()
+    def encode(self, tokenized_chunks: List[List[int]], maintain_state: bool = False) -> np.ndarray:
+        """
+        Genera embeddings para una lista de secuencias tokenizadas.
+        - maintain_state=True: Ingesta continua arrastrando la memoria Titan M entre chunks.
+        - maintain_state=False: Procesamiento independiente en paralelo.
+        """
+        if not tokenized_chunks:
+            return np.empty((0, self.embed_dim), dtype=np.float32)
+
+        # Validación DoS: Sanitización de entradas
+        for chunk in tokenized_chunks:
+            if not isinstance(chunk, list) or len(chunk) == 0:
+                raise ValueError("Cada chunk tokenizado debe ser una lista no vacía de enteros.")
+
+        device_type = 'cuda' if self.device.type == 'cuda' else 'cpu'
+        enable_autocast = (device_type == 'cuda') and not self.is_quantized
+        
+        with torch.autocast(device_type=device_type, dtype=torch.bfloat16, enabled=enable_autocast):
+            if maintain_state:
+                embeddings = []
+                M_state = self._current_memory_state
+                
+                for chunk in tokenized_chunks:
+                    input_tensor = torch.tensor([chunk], dtype=torch.long, device=self.device)
+                    emb, M_state = self.model(input_tensor, initial_state=M_state)
+                    embeddings.append(emb)
+                    
+                self._current_memory_state = M_state
+                embeddings = torch.cat(embeddings, dim=0)
+                return embeddings.float().cpu().numpy()
+            else:
+                max_len = max(len(chunk) for chunk in tokenized_chunks)
+                padded_inputs = [chunk + [0] * (max_len - len(chunk)) for chunk in tokenized_chunks]
+                input_tensor = torch.tensor(padded_inputs, dtype=torch.long, device=self.device)
+                embeddings, _ = self.model(input_tensor)
+                return embeddings.float().cpu().numpy()
+
+    # ===================================================================
+    # Interfaces Compatibles con LangChain / LlamaIndex / Vectorstores
+    # ===================================================================
+
+    def _tokenize_text(self, text: str, max_tokens: int = 512) -> List[int]:
+        """Tokenización segura con sanitización de longitud y filtrado de tokens de control."""
+        if not isinstance(text, str):
+            text = str(text)
+        output = self.tokenizer.encode(text)
+        ids = [t for t in output.ids if t not in (101, 102)]
+        return ids[:max_tokens] if ids else [0]
+
+    def embed_documents(self, texts: List[str], maintain_state: bool = False) -> List[List[float]]:
+        """
+        Interfaz LangChain/ChromaDB: Embebe una lista de textos documentales.
+        Usa tokenización multi-hilo en Rust.
+        """
+        if not texts:
+            return []
+            
+        batch_output = self.tokenizer.encode_batch([t if isinstance(t, str) else str(t) for t in texts])
+        tokenized = []
+        for enc in batch_output:
+            ids = [t for t in enc.ids if t not in (101, 102)][:512]
+            tokenized.append(ids if ids else [0])
+            
+        embeddings = self.encode(tokenized, maintain_state=maintain_state)
+        return embeddings.tolist()
+
+    def embed_query(self, text: str) -> List[float]:
+        """
+        Interfaz LangChain: Embebe una consulta de usuario.
+        Siempre opera en modo STATELESS para garantizar una búsqueda sin sesgos.
+        """
+        tokenized = self._tokenize_text(text)
+        embedding = self.encode([tokenized], maintain_state=False)
+        return embedding[0].tolist()
+
+    # ===================================================================
+    # Guardado y Carga de Hugging Face (Formatos Abiertos y Seguros)
+    # ===================================================================
+
+    def save_pretrained(self, save_directory: Union[str, Path]):
+        """
+        Exporta el modelo en el formato estándar de Hugging Face Hub:
+        - model.safetensors (pesos seguros)
+        - config.json (arquitectura y auto_map)
+        - tokenizer.json (tokenizador WordPiece en Rust)
+        """
+        save_dir = Path(save_directory)
+        save_dir.mkdir(parents=True, exist_ok=True)
+        
+        # 1. Guardar pesos en safetensors
+        from safetensors.torch import save_file
+        weights_path = save_dir / "model.safetensors"
+        save_file(self.model.state_dict(), str(weights_path))
+        
+        # 2. Guardar config.json
+        config_data = {
+            "architectures": ["TóarimEmbeddingModel"],
+            "model_type": "toarim",
+            "vocab_size": self.vocab_size,
+            "embed_dim": self.embed_dim,
+            "latent_k": self.latent_k,
+            "auto_map": {
+                "AutoModel": "toarim_embedder.TóarimEmbeddingModel"
+            }
+        }
+        with open(save_dir / "config.json", "w", encoding="utf-8") as f:
+            json.dump(config_data, f, indent=2)
+            
+        # 3. Guardar tokenizer.json
+        self.tokenizer.save(str(save_dir / "tokenizer.json"))
+        logger.info(f"Modelo Tóarim exportado exitosamente en: {save_dir}")
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, Path], **kwargs) -> "TóarimEmbedder":
+        """
+        Carga una instancia de TóarimEmbedder desde un directorio local o repositorio de HF.
+        """
+        model_dir = Path(pretrained_model_name_or_path)
+        
+        if model_dir.is_dir():
+            config_path = model_dir / "config.json"
+            weights_path = model_dir / "model.safetensors"
+            tokenizer_path = model_dir / "tokenizer.json"
+            
+            vocab_size = 30522
+            embed_dim = 768
+            latent_k = 128
+            
+            if config_path.exists():
+                with open(config_path, "r", encoding="utf-8") as f:
+                    cfg = json.load(f)
+                    vocab_size = cfg.get("vocab_size", vocab_size)
+                    embed_dim = cfg.get("embed_dim", embed_dim)
+                    latent_k = cfg.get("latent_k", latent_k)
+                    
+            return cls(
+                vocab_size=vocab_size,
+                embed_dim=embed_dim,
+                latent_k=latent_k,
+                model_path=str(weights_path) if weights_path.exists() else None,
+                tokenizer_name_or_path=str(tokenizer_path) if tokenizer_path.exists() else "bert-base-uncased",
+                **kwargs
+            )
+        else:
+            # Descarga remota desde Hugging Face Hub mediante huggingface_hub
+            from huggingface_hub import hf_hub_download
+            weights_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="model.safetensors")
+            config_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="config.json")
+            tokenizer_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="tokenizer.json")
+            
+            with open(config_file, "r", encoding="utf-8") as f:
+                cfg = json.load(f)
+                
+            return cls(
+                vocab_size=cfg.get("vocab_size", 30522),
+                embed_dim=cfg.get("embed_dim", 768),
+                latent_k=cfg.get("latent_k", 128),
+                model_path=weights_file,
+                tokenizer_name_or_path=tokenizer_file,
+                **kwargs
+            )
+
+
+# Alias ASCII estándar para máxima compatibilidad en importaciones
+ToarimEmbeddingModel = TóarimEmbeddingModel
+ToarimEmbedder = TóarimEmbedder
+ToarimEMbedder = TóarimEmbedder