ToarimEmbedder / toarim_embedder.py

Upload toarim_embedder.py with huggingface_hub

bef23b5 verified 5 days ago

22.8 kB

"""
=============================================================================
 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

	"""
	=============================================================================
	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