models/context_encoder.py

# models/context_encoder.py
import torch
import torch.nn as nn
from sentence_transformers import SentenceTransformer
from typing import List


class ContextEncoder(nn.Module):
    """
    Independent context encoder using Sentence-Transformer.
    Encodes session text into semantic embedding vectors.
    Frozen during training to preserve semantic space.
    """

    def __init__(
            self,
            model_name: str = "all-MiniLM-L6-v2",
            device: str = "cuda"
    ):
        super().__init__()

        self.device = device
        self.model_name = model_name

        # Load pre-trained Sentence-Transformer
        print(f"📥 Loading Context Encoder: {model_name}")
        self.encoder = SentenceTransformer(model_name, device=device)

        # Get embedding dimension
        self.embedding_dim = self.encoder.get_sentence_embedding_dimension()

        # Freeze all parameters (no gradient updates)
        for param in self.encoder.parameters():
            param.requires_grad = False

        print(f"✅ Context Encoder loaded. Embedding dim: {self.embedding_dim}")

    @torch.no_grad()
    def forward(self, texts: List[str]) -> torch.Tensor:
        """
        Encode list of text strings into embedding vectors (NO GRADIENT).
        Used for context encoding where gradient is not needed.

        Args:
            texts: List of text strings, length = batch_size

        Returns:
            embeddings: Tensor of shape (batch_size, embedding_dim)
        """
        embeddings = self.encoder.encode(
            texts,
            batch_size=len(texts),
            convert_to_tensor=True,
            show_progress_bar=False
        )

        embeddings = embeddings.to(self.device)

        return embeddings

    def encode_with_grad(self, texts: List[str]) -> torch.Tensor:
        """
        Encode texts WITH gradient tracking (for contrastive loss).
        Used for persona text encoding where we need gradient flow.

        Args:
            texts: List of text strings

        Returns:
            embeddings: Tensor of shape (batch_size, embedding_dim) with gradient tracking
        """
        # Tokenize manually
        encoded = self.encoder.tokenize(texts)

        # 🌟 FIX: Only move tensors, ignore strings/metadata
        encoded = {
            k: v.to(self.device)
            for k, v in encoded.items()
            if isinstance(v, torch.Tensor)
        }

        # Forward pass WITH gradient (no @torch.no_grad() decorator)
        sentence_features = self.encoder.forward(encoded)
        embeddings = sentence_features['sentence_embedding']

        # Normalize (same as default SentenceTransformer behavior)
        embeddings = embeddings / embeddings.norm(dim=1, keepdim=True)

        return embeddings

    def encode_batch(self, texts: List[str], batch_size: int = 32) -> torch.Tensor:
        """
        Encode large batch of texts with mini-batching to save memory.

        Args:
            texts: List of text strings
            batch_size: Mini-batch size for encoding

        Returns:
            embeddings: Tensor of shape (len(texts), embedding_dim)
        """
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            batch_embeddings = self.forward(batch_texts)
            all_embeddings.append(batch_embeddings)

        return torch.cat(all_embeddings, dim=0)


class ContextEncoderWrapper(nn.Module):
    """
    Wrapper for ContextEncoder to integrate with PyTorch training pipeline.
    Handles device placement and gradient flow.
    """

    def __init__(self, context_encoder: ContextEncoder):
        super().__init__()
        self.encoder = context_encoder

        # Move encoder to eval mode (no dropout, etc.)
        self.encoder.encoder.eval()

    def forward(self, texts: List[str]) -> torch.Tensor:
        """Forward pass with gradient disabled"""
        with torch.no_grad():
            return self.encoder(texts)

    def get_embedding_dim(self) -> int:
        """Get embedding dimension"""
        return self.encoder.embedding_dim