chatbot_model.py

import time
from transformers import TFAutoModel, AutoTokenizer
import tensorflow as tf
import numpy as np
from typing import List, Tuple, Dict, Optional, Union, Any
import math
from dataclasses import dataclass
import json
from pathlib import Path
import datetime       
import faiss
import gc
from tf_data_pipeline import TFDataPipeline
from response_quality_checker import ResponseQualityChecker
from cross_encoder_reranker import CrossEncoderReranker
from conversation_summarizer import DeviceAwareModel, Summarizer
from gpu_monitor import GPUMemoryMonitor
import absl.logging
from logger_config import config_logger
from tqdm.auto import tqdm

absl.logging.set_verbosity(absl.logging.WARNING)
logger = config_logger(__name__)

@dataclass
class ChatbotConfig:
    """Configuration for the RetrievalChatbot."""
    max_context_token_limit: int = 512
    embedding_dim: int = 768
    encoder_units: int = 256
    num_attention_heads: int = 8
    dropout_rate: float = 0.2
    l2_reg_weight: float = 0.001
    learning_rate: float = 0.001
    min_text_length: int = 3
    max_context_turns: int = 5
    warmup_steps: int = 200
    pretrained_model: str = 'distilbert-base-uncased'
    dtype: str = 'float32'
    freeze_embeddings: bool = False
    embedding_batch_size: int = 64
    search_batch_size: int = 64
    max_batch_size: int = 64
    neg_samples: int = 3
    max_retries: int = 3

    def to_dict(self) -> Dict:
        """Convert config to dictionary."""
        return {k: (str(v) if isinstance(v, Path) else v) 
                for k, v in self.__dict__.items()}

    @classmethod
    def from_dict(cls, config_dict: Dict) -> 'ChatbotConfig':
        """Create config from dictionary."""
        return cls(**{k: v for k, v in config_dict.items() 
                     if k in cls.__dataclass_fields__})

class EncoderModel(tf.keras.Model):
    """Dual encoder model with pretrained embeddings."""
    def __init__(
        self,
        config: ChatbotConfig,
        name: str = "encoder",
        **kwargs
    ):
        super().__init__(name=name, **kwargs)
        self.config = config

        # Load pretrained model
        self.pretrained = TFAutoModel.from_pretrained(config.pretrained_model)
        
        # Freeze layers based on config
        self._freeze_layers()
        
        # Pooling layer (Global Average Pooling)
        self.pooler = tf.keras.layers.GlobalAveragePooling1D()
        
        # Projection layer
        self.projection = tf.keras.layers.Dense(
            config.embedding_dim,
            activation='tanh',
            name="projection"
        )

        # Dropout and normalization
        self.dropout = tf.keras.layers.Dropout(config.dropout_rate)
        self.normalize = tf.keras.layers.Lambda(
            lambda x: tf.nn.l2_normalize(x, axis=1),
            name="l2_normalize"
        )

    def _freeze_layers(self):
        """Freeze layers of the pretrained model based on configuration."""
        if self.config.freeze_embeddings:
            self.pretrained.trainable = False
            logger.info("All pretrained layers frozen.")
        else:
            # Freeze only the first 'n' transformer layers
            for i, layer in enumerate(self.pretrained.layers):
                if isinstance(layer, tf.keras.layers.Layer):
                    if hasattr(layer, 'trainable'):
                        # Freeze the first transformer block
                        if i < 1:
                            layer.trainable = False
                            logger.info(f"Layer {i} frozen.")
                        else:
                            layer.trainable = True
                            
    def call(self, inputs: tf.Tensor, training: bool = False) -> tf.Tensor:
        """Forward pass."""
        # Get pretrained embeddings
        pretrained_outputs = self.pretrained(inputs, training=training)
        x = pretrained_outputs.last_hidden_state  # Shape: [batch_size, seq_len, embedding_dim]

        # Apply pooling, projection, dropout, and normalization
        x = self.pooler(x)  # Shape: [batch_size, 768]
        x = self.projection(x) # Shape: [batch_size, 768]
        x = self.dropout(x, training=training) # Apply dropout
        x = self.normalize(x)  # Shape: [batch_size, 768]

        return x

    def get_config(self) -> dict:
        """Return the config of the model."""
        config = super().get_config()
        config.update({
            "config": self.config.to_dict(),
            "name": self.name
        })
        return config
        
class RetrievalChatbot(DeviceAwareModel):
    """Retrieval-based chatbot using pretrained embeddings and FAISS for similarity search."""
    def __init__(
        self,
        config: ChatbotConfig,
        dialogues: List[dict] = [],
        device: str = None,
        strategy=None,
        reranker: Optional[CrossEncoderReranker] = None,
        summarizer: Optional[Summarizer] = None,
        mode: str = 'preparation' 
    ):
        super().__init__()
        self.config = config
        self.strategy = strategy
        self.device = device or self._setup_default_device()
        self.mode = mode.lower()
        
        # Initialize reranker, summarizer, tokenizer, and memory monitor
        self.reranker = reranker or self._initialize_reranker()
        self.summarizer = summarizer or self._initialize_summarizer()
        self.tokenizer = self._initialize_tokenizer()
        self.memory_monitor = GPUMemoryMonitor()
        
        # # Initialize models
        # self.min_batch_size = 8    
        # self.max_batch_size = 128
        # self.current_batch_size = 32
        
        # Initialize training history
        self.history = {
            "train_loss": [],
            "val_loss": [],
            "train_metrics": {},
            "val_metrics": {}
        }
        
        # Collect unique responses from dialogues
        if self.mode == 'preparation':
            # Collect unique responses from dialogues only in preparation mode
            self.response_pool, self.unique_responses = self._collect_responses(dialogues)
        else:
            # In training mode, assume response_pool is handled via TFRecord
            self.response_pool = []
            self.unique_responses = []
    
    def _setup_default_device(self) -> str:
        """Set up default device if none is provided."""
        if tf.config.list_physical_devices('GPU'):
            return 'GPU'
        else:
            return 'CPU'

    def _initialize_reranker(self) -> CrossEncoderReranker:
        """Initialize the CrossEncoderReranker."""
        logger.info("Initializing default CrossEncoderReranker...")
        return CrossEncoderReranker(model_name="cross-encoder/ms-marco-MiniLM-L-12-v2")

    def _initialize_summarizer(self) -> Summarizer:
        """Initialize the Summarizer."""
        logger.info("Initializing default Summarizer...")
        return Summarizer(device=self.device)

    def _initialize_tokenizer(self) -> AutoTokenizer:
        """Initialize the tokenizer and add special tokens."""
        logger.info("Initializing tokenizer and adding special tokens...")
        tokenizer = AutoTokenizer.from_pretrained(self.config.pretrained_model)
        special_tokens = {
            "user": "<USER>",
            "assistant": "<ASSISTANT>",
            "context": "<CONTEXT>",
            "sep": "<SEP>"
        }
        tokenizer.add_special_tokens(
            {'additional_special_tokens': list(special_tokens.values())}
        )
        return tokenizer
    
    def _collect_responses(self, dialogues: List[dict]) -> Tuple[List[str], List[str]]:
        """
        Collect unique responses from dialogues.
        Returns:
            response_pool: List of all possible responses.
            unique_responses: List of unique responses.
        """
        logger.info("Collecting unique responses from dialogues...")
        responses = set()
        for dialogue in dialogues:
            turns = dialogue.get('turns', [])
            for turn in turns:
                if turn.get('speaker') == 'assistant' and 'text' in turn:
                    response = turn['text'].strip()
                    if len(response) >= self.config.min_text_length:
                        responses.add(response)
        response_pool = list(responses)
        unique_responses = list(responses)  # Assuming uniqueness
        logger.info(f"Collected {len(response_pool)} unique responses.")
        return response_pool, unique_responses
    
    def build_models(self):
        """Initialize the shared encoder and FAISS index."""
        logger.info("Building encoder model...")
        tf.keras.backend.clear_session()
        
        # Shared encoder for both queries and responses
        self.encoder = EncoderModel(
            self.config,
            name="shared_encoder",
        )
        
        # Resize token embeddings after adding special tokens
        new_vocab_size = len(self.tokenizer)
        self.encoder.pretrained.resize_token_embeddings(new_vocab_size)
        logger.info(f"Token embeddings resized to: {new_vocab_size}")
        
        if self.mode == 'preparation':
            # Initialize FAISS index only in preparation mode
            self._initialize_faiss()
            # Compute and index embeddings
            self._compute_and_index_embeddings()
        else:
            # In training mode, skip FAISS indexing from dialogues
            logger.info("Training mode: Skipping FAISS index initialization from dialogues.")
        
        # Retrieve embedding dimension from encoder
        embedding_dim = self.config.embedding_dim 
        vocab_size = len(self.tokenizer)
        
        logger.info(f"Encoder Embedding Dimension: {embedding_dim}")
        logger.info(f"Encoder Embedding Vocabulary Size: {vocab_size}")
        if vocab_size >= embedding_dim:
            logger.info("Encoder model built and embeddings resized successfully.")
        else:
            logger.error("Vocabulary size is less than embedding dimension.")
            raise ValueError("Vocabulary size is less than embedding dimension.")
    
    def _adjust_batch_size(self) -> None:
        """Dynamically adjust batch size based on GPU memory usage."""
        if self.memory_monitor.should_reduce_batch_size():
            new_size = max(self.min_batch_size, self.current_batch_size // 2)
            if new_size != self.current_batch_size:
                logger.info(f"Reducing batch size to {new_size} due to high memory usage")
                self.current_batch_size = new_size
                gc.collect()
                if tf.config.list_physical_devices('GPU'):
                    tf.keras.backend.clear_session()
        elif self.memory_monitor.can_increase_batch_size():
            new_size = min(self.max_batch_size, self.current_batch_size * 2)
            if new_size != self.current_batch_size:
                logger.info(f"Increasing batch size to {new_size}")
                self.current_batch_size = new_size
            
    def _initialize_faiss(self):
        """Initialize FAISS with safe GPU handling and memory monitoring."""
        logger.info("Initializing FAISS index...")
        
        # Detect if we have GPU-enabled FAISS
        self.faiss_gpu = False
        self.gpu_resources = []
        
        try:
            if hasattr(faiss, 'get_num_gpus'):
                ngpus = faiss.get_num_gpus()
                if ngpus > 0:
                    # Configure GPU resources with memory limit
                    for i in range(ngpus):
                        res = faiss.StandardGpuResources()
                        # Set temp memory to 1/4 of total memory to avoid OOM
                        if self.memory_monitor.has_gpu:
                            stats = self.memory_monitor.get_memory_stats()
                            if stats:
                                temp_memory = int(stats.total * 0.25)  # 25% of total memory
                                res.setTempMemory(temp_memory)
                        self.gpu_resources.append(res)
                    self.faiss_gpu = True
                    logger.info(f"FAISS GPU resources initialized on {ngpus} GPUs")
        except Exception as e:
            logger.warning(f"Using CPU due to GPU initialization error: {e}")
        
        try:
            # Create appropriate index based on dataset size
            if len(self.unique_responses) < 1000:
                logger.info("Small dataset detected, using simple FlatIP index")
                self.index = faiss.IndexFlatIP(self.config.embedding_dim)
            else:
                # For larger datasets, consider using more efficient indices like IVF
                self.index = faiss.IndexFlatIP(self.config.embedding_dim)
            
            # Move to GPU(s) if available and needed
            if self.faiss_gpu and self.gpu_resources:
                try:
                    if len(self.gpu_resources) > 1:
                        self.index = faiss.index_cpu_to_gpus_list(self.index, self.gpu_resources)
                        logger.info("FAISS index distributed across multiple GPUs")
                    else:
                        self.index = faiss.index_cpu_to_gpu(self.gpu_resources[0], 0, self.index)
                        logger.info("FAISS index moved to single GPU")
                except Exception as e:
                    logger.warning(f"Failed to move index to GPU: {e}. Falling back to CPU")
                    self.faiss_gpu = False
        except Exception as e:
            logger.error(f"Error initializing FAISS: {e}")
            raise

    def encode_responses(
        self, 
        responses: List[str], 
        batch_size: int = 64
    ) -> tf.Tensor:
        """
        Encodes responses with more conservative memory management.
        """
        if not responses:
            logger.info("No responses to encode. Returning empty tensor.")
            return tf.constant([], dtype=tf.float32)
        
        all_embeddings = []
        self.current_batch_size = batch_size
        
        if self.memory_monitor.has_gpu:
            batch_size = 128

        total_processed = 0
        
        with tqdm(total=len(responses), desc="Encoding responses") as pbar:
            while total_processed < len(responses):
                # Monitor memory and adjust batch size
                if self.memory_monitor.has_gpu:
                    gpu_usage = self.memory_monitor.get_memory_usage()
                    if gpu_usage > 0.8:  # Over 80% usage
                        self.current_batch_size = max(128, self.current_batch_size // 2)
                        logger.info(f"High GPU memory usage ({gpu_usage:.1%}), reducing batch size to {self.current_batch_size}")
                        gc.collect()
                        tf.keras.backend.clear_session()
                
                # Get batch
                end_idx = min(total_processed + self.current_batch_size, len(responses))
                batch_texts = responses[total_processed:end_idx]
                
                try:
                    # Tokenize
                    encodings = self.tokenizer(
                        batch_texts,
                        padding='max_length',
                        truncation=True,
                        max_length=self.config.max_context_token_limit,
                        return_tensors='tf'
                    )

                    # Encode
                    embeddings_batch = self.encoder(encodings['input_ids'], training=False)
                    
                    # Cast to float32
                    if embeddings_batch.dtype != tf.float32:
                        embeddings_batch = tf.cast(embeddings_batch, tf.float32)

                    # Store
                    all_embeddings.append(embeddings_batch)
                    
                    # Update progress
                    batch_processed = len(batch_texts)
                    total_processed += batch_processed
                    
                    # Update progress bar
                    if self.memory_monitor.has_gpu:
                        gpu_usage = self.memory_monitor.get_memory_usage()
                        pbar.set_postfix({
                            'GPU mem': f'{gpu_usage:.1%}',
                            'batch_size': self.current_batch_size
                        })
                    pbar.update(batch_processed)
                    
                    # Memory cleanup every 1000 samples
                    if total_processed % 1000 == 0:
                        gc.collect()
                        if tf.config.list_physical_devices('GPU'):
                            tf.keras.backend.clear_session()

                except tf.errors.ResourceExhaustedError:
                    logger.warning("GPU memory exhausted during encoding, reducing batch size")
                    self.current_batch_size = max(8, self.current_batch_size // 2)
                    continue
                    
                except Exception as e:
                    logger.error(f"Error during encoding: {str(e)}")
                    raise

        # Concatenate results
        if not all_embeddings:
            logger.info("No embeddings were encoded. Returning empty tensor.")
            return tf.constant([], dtype=tf.float32)
        
        if len(all_embeddings) == 1:
            final_embeddings = all_embeddings[0]
        else:
            final_embeddings = tf.concat(all_embeddings, axis=0)
        
        return final_embeddings

    def _train_faiss_index(self, response_embeddings: np.ndarray) -> None:
        """Train FAISS index with better memory management and robust fallback mechanisms."""
        if self.index.is_trained:
            logger.info("Index already trained, skipping training phase")
            return

        logger.info("Starting FAISS index training...")
        
        try:
            # First attempt: Try training with smaller subset
            subset_size = min(5000, len(response_embeddings))  # Reduced from 10000
            logger.info(f"Using {subset_size} samples for initial training attempt")
            subset_idx = np.random.choice(len(response_embeddings), subset_size, replace=False)
            training_embeddings = response_embeddings[subset_idx].copy()  # Make a copy
            
            # Ensure contiguous memory layout
            training_embeddings = np.ascontiguousarray(training_embeddings)
                
            # Force cleanup before training
            gc.collect()
            if tf.config.list_physical_devices('GPU'):
                tf.keras.backend.clear_session()
                
            # Verify data properties
            logger.info(f"FAISS training data shape: {training_embeddings.shape}")
            logger.info(f"FAISS training data dtype: {training_embeddings.dtype}")
            
            logger.info("Starting initial training attempt...")
            self.index.train(training_embeddings)
            logger.info("Training completed successfully")
            
        except (RuntimeError, Exception) as e:
            logger.warning(f"Initial training attempt failed: {str(e)}")
            logger.info("Attempting fallback strategy...")
            
            try:
                # Move to CPU for more stable training
                if self.faiss_gpu:
                    logger.info("Moving index to CPU for fallback training")
                    cpu_index = faiss.index_gpu_to_cpu(self.index)
                else:
                    cpu_index = self.index
                
                # Create simpler index type if needed
                if isinstance(cpu_index, faiss.IndexIVFFlat):
                    logger.info("Creating simpler FlatL2 index for fallback")
                    cpu_index = faiss.IndexFlatL2(self.config.embedding_dim)
                
                # Use even smaller subset for CPU training
                subset_size = min(2000, len(response_embeddings))
                subset_idx = np.random.choice(len(response_embeddings), subset_size, replace=False)
                fallback_embeddings = response_embeddings[subset_idx].copy()
                
                # Ensure data is properly formatted
                if not fallback_embeddings.flags['C_CONTIGUOUS']:
                    fallback_embeddings = np.ascontiguousarray(fallback_embeddings)
                if fallback_embeddings.dtype != np.float32:
                    fallback_embeddings = fallback_embeddings.astype(np.float32)
                
                # Train on CPU
                logger.info("Training fallback index on CPU...")
                cpu_index.train(fallback_embeddings)
                
                # Move back to GPU if needed
                if self.faiss_gpu:
                    logger.info("Moving trained index back to GPU...")
                    if len(self.gpu_resources) > 1:
                        self.index = faiss.index_cpu_to_gpus_list(cpu_index, self.gpu_resources)
                    else:
                        self.index = faiss.index_cpu_to_gpu(self.gpu_resources[0], 0, cpu_index)
                else:
                    self.index = cpu_index
                    
                logger.info("Fallback training completed successfully")
                
            except Exception as e2:
                logger.error(f"Fallback training also failed: {str(e2)}")
                logger.warning("Creating basic brute-force index as last resort")
                
                try:
                    # Create basic brute-force index as last resort
                    dim = response_embeddings.shape[1]
                    basic_index = faiss.IndexFlatL2(dim)
                    
                    if self.faiss_gpu:
                        if len(self.gpu_resources) > 1:
                            self.index = faiss.index_cpu_to_gpus_list(basic_index, self.gpu_resources)
                        else:
                            self.index = faiss.index_cpu_to_gpu(self.gpu_resources[0], 0, basic_index)
                    else:
                        self.index = basic_index
                        
                    logger.info("Basic index created as fallback")
                    
                except Exception as e3:
                    logger.error(f"All training attempts failed: {str(e3)}")
                    raise RuntimeError("Unable to create working FAISS index")

    def _add_vectors_to_index(self, response_embeddings: np.ndarray) -> None:
        """Add vectors to FAISS index with enhanced memory management."""
        logger.info("Starting vector addition process...")
        
        # Even smaller batches
        initial_batch_size = 128
        min_batch_size = 32
        max_batch_size = 1024
        
        total_added = 0
        retry_count = 0
        max_retries = 5
        
        while total_added < len(response_embeddings):
            try:
                # Monitor memory
                if self.memory_monitor.has_gpu:
                    gpu_usage = self.memory_monitor.get_memory_usage()
                    #logger.info(f"GPU memory usage before batch: {gpu_usage:.1%}")
                    
                    # Force cleanup if memory usage is high
                    if gpu_usage > 0.7:  # Lower threshold to 70%
                        logger.info("High memory usage detected, forcing cleanup")
                        gc.collect()
                        tf.keras.backend.clear_session()
                
                # Get batch
                end_idx = min(total_added + initial_batch_size, len(response_embeddings))
                batch = response_embeddings[total_added:end_idx]
                
                # Add batch
                self.index.add(batch)
                
                # Update progress
                batch_size = len(batch)
                total_added += batch_size
                
                # Memory cleanup every few batches
                if total_added % (initial_batch_size * 5) == 0:
                    gc.collect()
                    if tf.config.list_physical_devices('GPU'):
                        tf.keras.backend.clear_session()
                
                # Gradually increase batch size
                if initial_batch_size < max_batch_size:
                    initial_batch_size = min(initial_batch_size + 25, max_batch_size)
                
            except Exception as e:
                logger.warning(f"Error adding batch: {str(e)}")
                retry_count += 1
                
                if retry_count > max_retries:
                    logger.error("Max retries exceeded.")
                    raise
                
                # Reduce batch size
                initial_batch_size = max(min_batch_size, initial_batch_size // 2)
                logger.info(f"Reducing batch size to {initial_batch_size} and retrying...")
                
                # Cleanup
                gc.collect()
                if tf.config.list_physical_devices('GPU'):
                    tf.keras.backend.clear_session()
                
                time.sleep(1)  # Brief pause before retry
        
        logger.info(f"Successfully added all {total_added} vectors to index")

    def _add_vectors_cpu_fallback(self, remaining_embeddings: np.ndarray, already_added: int = 0) -> None:
        """CPU fallback with extra safeguards and progress tracking."""
        logger.info(f"CPU Fallback: Adding {len(remaining_embeddings)} remaining vectors...")
        
        try:
            # Move index to CPU
            if self.faiss_gpu:
                logger.info("Moving index to CPU...")
                cpu_index = faiss.index_gpu_to_cpu(self.index)
            else:
                cpu_index = self.index
            
            # Add remaining vectors on CPU with very small batches
            batch_size = 128
            total_added = already_added
            
            for i in range(0, len(remaining_embeddings), batch_size):
                end_idx = min(i + batch_size, len(remaining_embeddings))
                batch = remaining_embeddings[i:end_idx]
                
                # Add batch
                cpu_index.add(batch)
                
                # Update progress
                total_added += len(batch)
                if i % (batch_size * 10) == 0:
                    logger.info(f"Added {total_added} vectors total "
                            f"({i}/{len(remaining_embeddings)} in current phase)")
                
                # Periodic cleanup
                if i % (batch_size * 20) == 0:
                    gc.collect()
            
            # Move back to GPU if needed
            if self.faiss_gpu:
                logger.info("Moving index back to GPU...")
                if len(self.gpu_resources) > 1:
                    self.index = faiss.index_cpu_to_gpus_list(cpu_index, self.gpu_resources)
                else:
                    self.index = faiss.index_cpu_to_gpu(self.gpu_resources[0], 0, cpu_index)
            else:
                self.index = cpu_index
                
            logger.info("CPU fallback completed successfully")
            
        except Exception as e:
            logger.error(f"Error during CPU fallback: {str(e)}")
            raise

    def _compute_and_index_embeddings(self):
        """Compute embeddings and build FAISS index with simpler handling."""
        logger.info("Computing embeddings and indexing with FAISS...")
        
        try:
            # Encode responses with memory monitoring
            logger.info("Encoding unique responses")
            response_embeddings = self.encode_responses(self.unique_responses)
            response_embeddings = response_embeddings.numpy()
            
            # Memory cleanup after encoding
            gc.collect()
            if tf.config.list_physical_devices('GPU'):
                tf.keras.backend.clear_session()
            
            # Ensure float32 and memory contiguous
            response_embeddings = response_embeddings.astype('float32')
            response_embeddings = np.ascontiguousarray(response_embeddings)
            
            # Log memory state before normalization
            if self.memory_monitor.has_gpu:
                stats = self.memory_monitor.get_memory_stats()
                if stats:
                    logger.info(f"GPU memory before normalization: {stats.used/1e9:.2f}GB used")
            
            # Normalize embeddings
            logger.info("Normalizing embeddings with FAISS")
            faiss.normalize_L2(response_embeddings)
            
            # Create and initialize simple FlatIP index
            dim = response_embeddings.shape[1]
            if self.faiss_gpu:
                cpu_index = faiss.IndexFlatIP(dim)
                if len(self.gpu_resources) > 1:
                    self.index = faiss.index_cpu_to_gpus_list(cpu_index, self.gpu_resources)
                else:
                    self.index = faiss.index_cpu_to_gpu(self.gpu_resources[0], 0, cpu_index)
            else:
                self.index = faiss.IndexFlatIP(dim)
            
            # Add vectors to index
            self._add_vectors_to_index(response_embeddings)
            
            # Store responses and embeddings
            self.response_pool = self.unique_responses
            self.response_embeddings = response_embeddings
            
            # Final memory cleanup
            gc.collect()
            if tf.config.list_physical_devices('GPU'):
                tf.keras.backend.clear_session()
            
            # Log final state
            logger.info(f"Successfully indexed {self.index.ntotal} responses")
            if self.memory_monitor.has_gpu:
                stats = self.memory_monitor.get_memory_stats()
                if stats:
                    logger.info(f"Final GPU memory usage: {stats.used/1e9:.2f}GB used")
            
            logger.info("Indexing completed successfully")
            
        except Exception as e:
            logger.error(f"Error during indexing: {e}")
            # Ensure cleanup even on error
            gc.collect()
            if tf.config.list_physical_devices('GPU'):
                tf.keras.backend.clear_session()
            raise

    def verify_faiss_index(self):
        """Verify that FAISS index matches the response pool, if index exists."""
        if not hasattr(self, 'index') or self.index is None:
            logger.info("FAISS index not initialized. Skipping verification.")
            return
        
        indexed_size = self.index.ntotal
        pool_size = len(self.response_pool)
        logger.info(f"FAISS index size: {indexed_size}")
        logger.info(f"Response pool size: {pool_size}")
        if indexed_size != pool_size:
            logger.warning("Mismatch between FAISS index size and response pool size.")
        else:
            logger.info("FAISS index correctly matches the response pool.")

    def encode_query(self, query: str, context: Optional[List[Tuple[str, str]]] = None) -> tf.Tensor:
        """Encode a query with optional conversation context."""
        # Prepare query with context
        if context:
            context_str = ' '.join([
                f"{self.special_tokens['user']} {q} "
                f"{self.special_tokens['assistant']} {r}"
                for q, r in context[-self.config.max_context_turns:]
            ])
            query = f"{context_str} {self.special_tokens['user']} {query}"
        else:
            query = f"{self.special_tokens['user']} {query}"
        
        # Tokenize and encode
        encodings = self.tokenizer(
            [query],
            padding='max_length',
            truncation=True,
            max_length=self.config.max_context_token_limit,
            return_tensors='tf'
        )
        input_ids = encodings['input_ids']
        
        # Verify token IDs
        max_id = tf.reduce_max(input_ids).numpy()
        new_vocab_size = len(self.tokenizer)
        
        if max_id >= new_vocab_size:
            logger.error(f"Token ID {max_id} exceeds the vocabulary size {new_vocab_size}.")
            raise ValueError("Token ID exceeds vocabulary size.")
        
        # Get embeddings from the shared encoder
        return self.encoder(input_ids, training=False)

    def retrieve_responses_cross_encoder(
        self,
        query: str,
        top_k: int,
        reranker: Optional[CrossEncoderReranker] = None,
        summarizer: Optional[Summarizer] = None,
        summarize_threshold: int = 512  # Summarize over 512 tokens
    ) -> List[Tuple[str, float]]:
        """
        Retrieve top-k from FAISS, then re-rank them with a cross-encoder.
        Optionally summarize the user query if it's too long.
        """
        if reranker is None:
            reranker = self.reranker
        if summarizer is None:
            summarizer = self.summarizer
            
        # Optional summarization
        if summarizer and len(query.split()) > summarize_threshold:
            logger.info(f"Query is long. Summarizing before cross-encoder. Original length: {len(query.split())}")
            query = summarizer.summarize_text(query)
            logger.info(f"Summarized query: {query}")

        # 2) Dense retrieval
        dense_topk = self.retrieve_responses_faiss(query, top_k=top_k)  # [(resp, dense_score), ...]

        if not dense_topk:
            return []

        # 3) Cross-encoder rerank
        candidate_texts = [pair[0] for pair in dense_topk]
        cross_scores = reranker.rerank(query, candidate_texts, max_length=256)

        # Combine
        combined = [(text, score) for (text, _), score in zip(dense_topk, cross_scores)]
        # Sort descending by cross-encoder score
        combined.sort(key=lambda x: x[1], reverse=True)

        return combined
    
    def retrieve_responses_faiss(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
        """Retrieve top-k responses using FAISS."""
        if not hasattr(self, 'index') or self.index is None:
            logger.warning("FAISS index not initialized. Cannot retrieve responses.")
            return []
        
        # Encode the query
        q_emb = self.encode_query(query)  # Shape: [1, embedding_dim]
        q_emb_np = q_emb.numpy().astype('float32')  # Ensure type match
        
        # Normalize the query embedding for cosine similarity
        faiss.normalize_L2(q_emb_np)
        
        # Search the FAISS index
        distances, indices = self.index.search(q_emb_np, top_k)
        
        # Map indices to responses and distances to similarities
        top_responses = []
        for i, idx in enumerate(indices[0]):
            if idx < len(self.response_pool):
                top_responses.append((self.response_pool[idx], float(distances[0][i])))
            else:
                logger.warning(f"FAISS returned invalid index {idx}. Skipping.")
        
        return top_responses
    
    def save_models(self, save_dir: Union[str, Path]):
        """Save models and configuration."""
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        
        # Save config
        with open(save_dir / "config.json", "w") as f:
            json.dump(self.config.to_dict(), f, indent=2)
        
        # Save models
        self.encoder.pretrained.save_pretrained(save_dir / "shared_encoder")
        
        # Save tokenizer
        self.tokenizer.save_pretrained(save_dir / "tokenizer")
        
        logger.info(f"Models and tokenizer saved to {save_dir}.")
    
    @classmethod
    def load_models(cls, load_dir: Union[str, Path]) -> 'RetrievalChatbot':
        """Load saved models and configuration."""
        load_dir = Path(load_dir)
        
        # Load config
        with open(load_dir / "config.json", "r") as f:
            config = ChatbotConfig.from_dict(json.load(f))
        
        # Initialize chatbot
        chatbot = cls(config)
        
        # Load models
        chatbot.encoder.pretrained = TFAutoModel.from_pretrained(
            load_dir / "shared_encoder",
            config=config
        )
        
        # Load tokenizer
        chatbot.tokenizer = AutoTokenizer.from_pretrained(load_dir / "tokenizer")
        
        logger.info(f"Models and tokenizer loaded from {load_dir}.")
        return chatbot
    
    def train_streaming(
        self,
        tfrecord_file_path: str,
        epochs: int = 20,
        batch_size: int = 16,
        validation_split: float = 0.2,
        checkpoint_dir: str = "checkpoints/",
        use_lr_schedule: bool = True,
        peak_lr: float = 2e-5,
        warmup_steps_ratio: float = 0.1,
        early_stopping_patience: int = 3,
        min_delta: float = 1e-4,
    ) -> None:
        """Training using a pre-prepared TFRecord dataset."""
        logger.info("Starting training with pre-prepared TFRecord dataset...")
        
        def parse_tfrecord_fn(example_proto, max_length, neg_samples):
            """
            Parses a single TFRecord example.

            Args:
                example_proto: A serialized TFRecord example.
                max_length: The maximum sequence length for tokenization.
                neg_samples: The number of hard negatives per query.

            Returns:
                A tuple of (query_ids, positive_ids, negative_ids).
            """
            feature_description = {
                'query_ids': tf.io.FixedLenFeature([max_length], tf.int64),
                'positive_ids': tf.io.FixedLenFeature([max_length], tf.int64),
                'negative_ids': tf.io.FixedLenFeature([neg_samples * max_length], tf.int64),
            }
            parsed_features = tf.io.parse_single_example(example_proto, feature_description)

            query_ids = tf.cast(parsed_features['query_ids'], tf.int32)
            positive_ids = tf.cast(parsed_features['positive_ids'], tf.int32)
            negative_ids = tf.cast(parsed_features['negative_ids'], tf.int32)
            negative_ids = tf.reshape(negative_ids, [neg_samples, max_length])

            return query_ids, positive_ids, negative_ids

        # Calculate total steps by counting the number of records in the TFRecord
        raw_dataset = tf.data.TFRecordDataset(tfrecord_file_path)
        total_pairs = sum(1 for _ in raw_dataset)
        logger.info(f"Total pairs in TFRecord: {total_pairs}")

        train_size = int(total_pairs * (1 - validation_split))
        val_size = total_pairs - train_size
        steps_per_epoch = math.ceil(train_size / batch_size)
        val_steps = math.ceil(val_size / batch_size)
        total_steps = steps_per_epoch * epochs

        logger.info(f"Training pairs: {train_size}")
        logger.info(f"Validation pairs: {val_size}")
        logger.info(f"Steps per epoch: {steps_per_epoch}")
        logger.info(f"Validation steps: {val_steps}")
        logger.info(f"Total steps: {total_steps}")

        # Set up optimizer with learning rate schedule
        if use_lr_schedule:
            warmup_steps = int(total_steps * warmup_steps_ratio)
            lr_schedule = self._get_lr_schedule(
                total_steps=total_steps,
                peak_lr=peak_lr,
                warmup_steps=warmup_steps
            )
            self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
            logger.info("Using custom learning rate schedule.")
        else:
            self.optimizer = tf.keras.optimizers.Adam(learning_rate=peak_lr)
            logger.info("Using fixed learning rate.")

        # Initialize checkpoint manager
        checkpoint = tf.train.Checkpoint(optimizer=self.optimizer, model=self.encoder)
        manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)

        # Setup TensorBoard
        log_dir = Path(checkpoint_dir) / "tensorboard_logs"
        log_dir.mkdir(parents=True, exist_ok=True)
        current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
        train_log_dir = str(log_dir / f"train_{current_time}")
        val_log_dir = str(log_dir / f"val_{current_time}")
        train_summary_writer = tf.summary.create_file_writer(train_log_dir)
        val_summary_writer = tf.summary.create_file_writer(val_log_dir)
        logger.info(f"TensorBoard logs will be saved in {log_dir}")

        # Define the parsing function with the appropriate max_length and neg_samples
        parse_fn = lambda x: parse_tfrecord_fn(x, self.config.max_context_token_limit, self.config.neg_samples)

        # Create the full dataset
        dataset = tf.data.TFRecordDataset(tfrecord_file_path)
        dataset = dataset.map(parse_fn, num_parallel_calls=tf.data.AUTOTUNE)
        dataset = dataset.shuffle(buffer_size=10000)  # Adjust buffer size as needed
        dataset = dataset.batch(batch_size, drop_remainder=True)
        dataset = dataset.prefetch(tf.data.AUTOTUNE)

        # Split into training and validation
        train_dataset = dataset.take(train_size)
        val_dataset = dataset.skip(train_size).take(val_size)

        # Training loop
        best_val_loss = float("inf")
        epochs_no_improve = 0

        for epoch in range(1, epochs + 1):
            # --- Training Phase ---
            epoch_loss_avg = tf.keras.metrics.Mean()
            batches_processed = 0

            try:
                train_pbar = tqdm(total=steps_per_epoch, desc=f"Training Epoch {epoch}", unit="batch")
                is_tqdm_train = True
            except ImportError:
                train_pbar = None
                is_tqdm_train = False
                logger.info("Training progress bar disabled")

            for q_batch, p_batch, n_batch in train_dataset:
                loss = self.train_step(q_batch, p_batch, n_batch)
                epoch_loss_avg(loss)
                batches_processed += 1

                # Log to TensorBoard
                with train_summary_writer.as_default():
                    tf.summary.scalar("loss", loss, step=(epoch - 1) * steps_per_epoch + batches_processed)

                # Update progress bar
                if use_lr_schedule:
                    current_lr = float(lr_schedule(self.optimizer.iterations))
                else:
                    current_lr = float(self.optimizer.learning_rate.numpy())

                if is_tqdm_train:
                    train_pbar.update(1)
                    train_pbar.set_postfix({
                        "loss": f"{loss.numpy():.4f}",
                        "lr": f"{current_lr:.2e}",
                        "batches": f"{batches_processed}/{steps_per_epoch}"
                    })

                # Memory cleanup
                gc.collect()

                if batches_processed >= steps_per_epoch:
                    break

            if is_tqdm_train and train_pbar:
                train_pbar.close()

            # --- Validation Phase ---
            val_loss_avg = tf.keras.metrics.Mean()
            val_batches_processed = 0

            try:
                val_pbar = tqdm(total=val_steps, desc="Validation", unit="batch")
                is_tqdm_val = True
            except ImportError:
                val_pbar = None
                is_tqdm_val = False
                logger.info("Validation progress bar disabled")

            for q_batch, p_batch, n_batch in val_dataset:
                val_loss = self.validation_step(q_batch, p_batch, n_batch)
                val_loss_avg(val_loss)
                val_batches_processed += 1

                if is_tqdm_val:
                    val_pbar.update(1)
                    val_pbar.set_postfix({
                        "val_loss": f"{val_loss.numpy():.4f}",
                        "batches": f"{val_batches_processed}/{val_steps}"
                    })

                # Memory cleanup
                gc.collect()

                if val_batches_processed >= val_steps:
                    break

            if is_tqdm_val and val_pbar:
                val_pbar.close()

            # End of epoch: compute final epoch stats, log, and save checkpoint
            train_loss = epoch_loss_avg.result().numpy()
            val_loss = val_loss_avg.result().numpy()
            logger.info(f"Epoch {epoch} Complete: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}")

            # Log epoch metrics
            with train_summary_writer.as_default():
                tf.summary.scalar("epoch_loss", train_loss, step=epoch)
            with val_summary_writer.as_default():
                tf.summary.scalar("val_loss", val_loss, step=epoch)

            # Save checkpoint
            manager.save()

            # Store metrics in history
            self.history['train_loss'].append(train_loss)
            self.history['val_loss'].append(val_loss)

            if use_lr_schedule:
                current_lr = float(lr_schedule(self.optimizer.iterations))
            else:
                current_lr = float(self.optimizer.learning_rate.numpy())

            self.history.setdefault('learning_rate', []).append(current_lr)

            # Early stopping logic
            if val_loss < best_val_loss - min_delta:
                best_val_loss = val_loss
                epochs_no_improve = 0
                logger.info(f"Validation loss improved to {val_loss:.4f}. Reset patience.")
            else:
                epochs_no_improve += 1
                logger.info(f"No improvement this epoch. Patience: {epochs_no_improve}/{early_stopping_patience}")
                if epochs_no_improve >= early_stopping_patience:
                    logger.info("Early stopping triggered.")
                    break

        logger.info("Training completed!")
        
    @tf.function
    def train_step(
        self, 
        q_batch: tf.Tensor, 
        p_batch: tf.Tensor, 
        n_batch: tf.Tensor
    ) -> tf.Tensor:
        """
        Single training step using queries, positives, and hard negatives.
        """
        with tf.GradientTape() as tape:
            # Encode queries
            q_enc = self.encoder(q_batch, training=True)  # [batch_size, embed_dim]

            # Encode positives
            p_enc = self.encoder(p_batch, training=True)  # [batch_size, embed_dim]

            # Encode negatives
            # n_batch: [batch_size, neg_samples, max_length]
            shape = tf.shape(n_batch)
            bs = shape[0]
            neg_samples = shape[1]

            # Flatten negatives to feed them in one pass:
            # => [batch_size * neg_samples, max_length]
            n_batch_flat = tf.reshape(n_batch, [bs * neg_samples, shape[2]])
            n_enc_flat = self.encoder(n_batch_flat, training=True)  # [bs*neg_samples, embed_dim]

            # Reshape back => [batch_size, neg_samples, embed_dim]
            n_enc = tf.reshape(n_enc_flat, [bs, neg_samples, -1])

            # Combine the positive embedding and negative embeddings along dim=1
            # => shape [batch_size, 1 + neg_samples, embed_dim]
            # The first column is the positive; subsequent columns are negatives
            combined_p_n = tf.concat(
                [tf.expand_dims(p_enc, axis=1), n_enc], 
                axis=1
            )  # [bs, (1+neg_samples), embed_dim]

            # Now compute scores: dot product of q_enc with each column in combined_p_n
            # We'll use `tf.einsum` to handle the batch dimension properly
            # dot_products => shape [batch_size, (1+neg_samples)]
            dot_products = tf.einsum('bd,bkd->bk', q_enc, combined_p_n)

            # The label for each row is 0 (the first column is the correct/positive)
            labels = tf.zeros([bs], dtype=tf.int32)

            # Cross-entropy over the [batch_size, 1+neg_samples] scores
            loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=labels,
                logits=dot_products
            )
            loss = tf.reduce_mean(loss)

        # Apply gradients
        gradients = tape.gradient(loss, self.encoder.trainable_variables)
        self.optimizer.apply_gradients(zip(gradients, self.encoder.trainable_variables))
        return loss

    @tf.function
    def validation_step(
        self, 
        q_batch: tf.Tensor, 
        p_batch: tf.Tensor, 
        n_batch: tf.Tensor
    ) -> tf.Tensor:
        """
        Single validation step using queries, positives, and hard negatives.
        """
        q_enc = self.encoder(q_batch, training=False)
        p_enc = self.encoder(p_batch, training=False)

        shape = tf.shape(n_batch)
        bs = shape[0]
        neg_samples = shape[1]

        n_batch_flat = tf.reshape(n_batch, [bs * neg_samples, shape[2]])
        n_enc_flat = self.encoder(n_batch_flat, training=False)
        n_enc = tf.reshape(n_enc_flat, [bs, neg_samples, -1])

        combined_p_n = tf.concat(
            [tf.expand_dims(p_enc, axis=1), n_enc],
            axis=1
        )

        dot_products = tf.einsum('bd,bkd->bk', q_enc, combined_p_n)
        labels = tf.zeros([bs], dtype=tf.int32)

        loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=labels,
            logits=dot_products
        )
        loss = tf.reduce_mean(loss)

        return loss
    # def train_streaming(
    #     self,
    #     dialogues: List[dict],
    #     epochs: int = 20,
    #     batch_size: int = 16,
    #     validation_split: float = 0.2,
    #     checkpoint_dir: str = "checkpoints/",
    #     use_lr_schedule: bool = True,
    #     peak_lr: float = 2e-5,
    #     warmup_steps_ratio: float = 0.1,
    #     early_stopping_patience: int = 3,
    #     min_delta: float = 1e-4,
    #     neg_samples: int = 1
    # ) -> None:
    #     """Streaming training with tf.data pipeline."""
    #     logger.info("Starting streaming training pipeline with tf.data...")

    #     # Initialize TFDataPipeline (replaces StreamingDataPipeline)
    #     dataset_preparer = TFDataPipeline(
    #         embedding_batch_size=self.config.embedding_batch_size,
    #         tokenizer=self.tokenizer,
    #         encoder=self.encoder,
    #         index=self.index,  # Pass CPU version of FAISS index
    #         response_pool=self.response_pool,
    #         max_length=self.config.max_context_token_limit,
    #         neg_samples=neg_samples
    #     )

    #     # Calculate total steps for learning rate schedule
    #     total_pairs = dataset_preparer.estimate_total_pairs(dialogues)
    #     train_size = int(total_pairs * (1 - validation_split))
    #     val_size = int(total_pairs * validation_split)
    #     steps_per_epoch = int(math.ceil(train_size / batch_size))
    #     val_steps = int(math.ceil(val_size / batch_size))
    #     total_steps = steps_per_epoch * epochs

    #     logger.info(f"Total pairs: {total_pairs}")
    #     logger.info(f"Training pairs: {train_size}")
    #     logger.info(f"Validation pairs: {val_size}")
    #     logger.info(f"Steps per epoch: {steps_per_epoch}")
    #     logger.info(f"Validation steps: {val_steps}")
    #     logger.info(f"Total steps: {total_steps}")

    #     # Set up optimizer with learning rate schedule
    #     if use_lr_schedule:
    #         warmup_steps = int(total_steps * warmup_steps_ratio)
    #         lr_schedule = self._get_lr_schedule(
    #             total_steps=total_steps,
    #             peak_lr=peak_lr,
    #             warmup_steps=warmup_steps
    #         )
    #         self.optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
    #         logger.info("Using custom learning rate schedule.")
    #     else:
    #         self.optimizer = tf.keras.optimizers.Adam(learning_rate=peak_lr)
    #         logger.info("Using fixed learning rate.")

    #     # Initialize checkpoint manager
    #     checkpoint = tf.train.Checkpoint(optimizer=self.optimizer, model=self.encoder)
    #     manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)

    #     # Setup TensorBoard
    #     log_dir = Path(checkpoint_dir) / "tensorboard_logs"
    #     log_dir.mkdir(parents=True, exist_ok=True)
    #     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
    #     train_log_dir = str(log_dir / f"train_{current_time}")
    #     val_log_dir = str(log_dir / f"val_{current_time}")
    #     train_summary_writer = tf.summary.create_file_writer(train_log_dir)
    #     val_summary_writer = tf.summary.create_file_writer(val_log_dir)
    #     logger.info(f"TensorBoard logs will be saved in {log_dir}")

    #     # Create training and validation datasets
    #     train_dataset = dataset_preparer.get_tf_dataset(dialogues, batch_size).take(train_size)
    #     val_dataset = dataset_preparer.get_tf_dataset(dialogues, batch_size).skip(train_size).take(val_size)

    #     # Training loop
    #     best_val_loss = float("inf")
    #     epochs_no_improve = 0

    #     for epoch in range(1, epochs + 1):
    #         # --- Training Phase ---
    #         epoch_loss_avg = tf.keras.metrics.Mean()
    #         batches_processed = 0

    #         try:
    #             train_pbar = tqdm(total=steps_per_epoch, desc=f"Training Epoch {epoch}", unit="batch")
    #             is_tqdm_train = True
    #         except ImportError:
    #             train_pbar = None
    #             is_tqdm_train = False
    #             logger.info("Training progress bar disabled")

    #         for q_batch, p_batch, n_batch in train_dataset:
    #             #p_batch = p_n_batch[:, 0, :] # Extract positive from (positive, negative) pair
    #             loss = self.train_step(q_batch, p_batch, n_batch)
    #             epoch_loss_avg(loss)
    #             batches_processed += 1

    #             # Log to TensorBoard
    #             with train_summary_writer.as_default():
    #                 tf.summary.scalar("loss", loss, step=(epoch - 1) * steps_per_epoch + batches_processed)

    #             # Update progress bar
    #             if use_lr_schedule:
    #                 current_lr = float(lr_schedule(self.optimizer.iterations))
    #             else:
    #                 current_lr = float(self.optimizer.learning_rate.numpy())

    #             if is_tqdm_train:
    #                 train_pbar.update(1)
    #                 train_pbar.set_postfix({
    #                     "loss": f"{loss.numpy():.4f}",
    #                     "lr": f"{current_lr:.2e}",
    #                     "batches": f"{batches_processed}/{steps_per_epoch}"
    #                 })
                
    #             # Memory cleanup
    #             gc.collect()

    #             if batches_processed >= steps_per_epoch:
    #                 break

    #         if is_tqdm_train and train_pbar:
    #             train_pbar.close()

    #         # --- Validation Phase ---
    #         val_loss_avg = tf.keras.metrics.Mean()
    #         val_batches_processed = 0

    #         try:
    #             val_pbar = tqdm(total=val_steps, desc="Validation", unit="batch")
    #             is_tqdm_val = True
    #         except ImportError:
    #             val_pbar = None
    #             is_tqdm_val = False
    #             logger.info("Validation progress bar disabled")

    #         for q_batch, p_batch, n_batch in val_dataset:
    #             #p_batch = p_n_batch[:, 0, :] # Extract positive from (positive, negative) pair
    #             val_loss = self.validation_step(q_batch, p_batch, n_batch)
    #             val_loss_avg(val_loss)
    #             val_batches_processed += 1

    #             if is_tqdm_val:
    #                 val_pbar.update(1)
    #                 val_pbar.set_postfix({
    #                     "val_loss": f"{val_loss.numpy():.4f}",
    #                     "batches": f"{val_batches_processed}/{val_steps}"
    #                 })
                    
    #             # Memory cleanup
    #             gc.collect()


    #             if val_batches_processed >= val_steps:
    #                 break

    #         if is_tqdm_val and val_pbar:
    #             val_pbar.close()

    #         # End of epoch: compute final epoch stats, log, and save checkpoint
    #         train_loss = epoch_loss_avg.result().numpy()
    #         val_loss = val_loss_avg.result().numpy()
    #         logger.info(f"Epoch {epoch} Complete: Train Loss={train_loss:.4f}, Val Loss={val_loss:.4f}")

    #         # Log epoch metrics
    #         with train_summary_writer.as_default():
    #             tf.summary.scalar("epoch_loss", train_loss, step=epoch)
    #         with val_summary_writer.as_default():
    #             tf.summary.scalar("val_loss", val_loss, step=epoch)

    #         # Save checkpoint
    #         manager.save()

    #         # Store metrics in history
    #         self.history['train_loss'].append(train_loss)
    #         self.history['val_loss'].append(val_loss)

    #         if use_lr_schedule:
    #             current_lr = float(lr_schedule(self.optimizer.iterations))
    #         else:
    #             current_lr = float(self.optimizer.learning_rate.numpy())

    #         self.history.setdefault('learning_rate', []).append(current_lr)

    #         # Early stopping logic
    #         if val_loss < best_val_loss - min_delta:
    #             best_val_loss = val_loss
    #             epochs_no_improve = 0
    #             logger.info(f"Validation loss improved to {val_loss:.4f}. Reset patience.")
    #         else:
    #             epochs_no_improve += 1
    #             logger.info(f"No improvement this epoch. Patience: {epochs_no_improve}/{early_stopping_patience}")
    #             if epochs_no_improve >= early_stopping_patience:
    #                 logger.info("Early stopping triggered.")
    #                 break

    #     logger.info("Streaming training completed!")


    # @tf.function
    # def train_step(
    #     self, 
    #     q_batch: tf.Tensor, 
    #     p_batch: tf.Tensor, 
    #     n_batch: tf.Tensor, 
    #     attention_mask: Optional[tf.Tensor] = None
    # ) -> tf.Tensor:
    #     """
    #     Single training step that uses queries, positives, and negatives in a 
    #     contrastive/InfoNCE style. The label is always 0 (the positive) vs. 
    #     the negative alternatives.
    #     """
    #     with tf.GradientTape() as tape:
    #         # Encode queries
    #         q_enc = self.encoder(q_batch, training=True)  # [batch_size, embed_dim]

    #         # Encode positives
    #         p_enc = self.encoder(p_batch, training=True)  # [batch_size, embed_dim]

    #         # Encode negatives
    #         # n_batch: [batch_size, neg_samples, max_length]
    #         shape = tf.shape(n_batch)
    #         bs = shape[0]
    #         neg_samples = shape[1]

    #         # Flatten negatives to feed them in one pass:
    #         # => [batch_size * neg_samples, max_length]
    #         n_batch_flat = tf.reshape(n_batch, [bs * neg_samples, shape[2]])
    #         n_enc_flat = self.encoder(n_batch_flat, training=True)  # [bs*neg_samples, embed_dim]

    #         # Reshape back => [batch_size, neg_samples, embed_dim]
    #         n_enc = tf.reshape(n_enc_flat, [bs, neg_samples, -1])

    #         # Combine the positive embedding and negative embeddings along dim=1
    #         # => shape [batch_size, 1 + neg_samples, embed_dim]
    #         # The first column is the positive; subsequent columns are negatives
    #         combined_p_n = tf.concat(
    #             [tf.expand_dims(p_enc, axis=1), n_enc], 
    #             axis=1
    #         )  # [bs, (1+neg_samples), embed_dim]

    #         # Now compute scores: dot product of q_enc with each column in combined_p_n
    #         # We'll use `tf.einsum` to handle the batch dimension properly
    #         # dot_products => shape [batch_size, (1+neg_samples)]
    #         dot_products = tf.einsum('bd,bkd->bk', q_enc, combined_p_n)

    #         # The label for each row is 0 (the first column is the correct/positive)
    #         labels = tf.zeros([bs], dtype=tf.int32)

    #         # Cross-entropy over the [batch_size, 1+neg_samples] scores
    #         loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
    #             labels=labels,
    #             logits=dot_products
    #         )
    #         loss = tf.reduce_mean(loss)

    #         # If there's an attention_mask you want to apply (less common in this scenario),
    #         # you could do something like:
    #         if attention_mask is not None:
    #             loss = loss * attention_mask
    #             loss = tf.reduce_sum(loss) / tf.reduce_sum(attention_mask)

    #     # Apply gradients
    #     gradients = tape.gradient(loss, self.encoder.trainable_variables)
    #     self.optimizer.apply_gradients(zip(gradients, self.encoder.trainable_variables))
    #     return loss

    # @tf.function
    # def validation_step(
    #     self, 
    #     q_batch: tf.Tensor, 
    #     p_batch: tf.Tensor, 
    #     n_batch: tf.Tensor, 
    #     attention_mask: Optional[tf.Tensor] = None
    # ) -> tf.Tensor:
    #     """
    #     Single validation step with queries, positives, and negatives. 
    #     Uses the same loss calculation as train_step, but `training=False`.
    #     """
    #     q_enc = self.encoder(q_batch, training=False)
    #     p_enc = self.encoder(p_batch, training=False)

    #     shape = tf.shape(n_batch)
    #     bs = shape[0]
    #     neg_samples = shape[1]

    #     n_batch_flat = tf.reshape(n_batch, [bs * neg_samples, shape[2]])
    #     n_enc_flat = self.encoder(n_batch_flat, training=False)
    #     n_enc = tf.reshape(n_enc_flat, [bs, neg_samples, -1])

    #     combined_p_n = tf.concat(
    #         [tf.expand_dims(p_enc, axis=1), n_enc],
    #         axis=1
    #     )

    #     dot_products = tf.einsum('bd,bkd->bk', q_enc, combined_p_n)
    #     labels = tf.zeros([bs], dtype=tf.int32)

    #     loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
    #         labels=labels,
    #         logits=dot_products
    #     )
    #     loss = tf.reduce_mean(loss)

    #     if attention_mask is not None:
    #         loss = loss * attention_mask
    #         loss = tf.reduce_sum(loss) / tf.reduce_sum(attention_mask)

    #     return loss
    
    def _get_lr_schedule(
        self,
        total_steps: int,
        peak_lr: float,
        warmup_steps: int
    ) -> tf.keras.optimizers.schedules.LearningRateSchedule:
        """Create a custom learning rate schedule with warmup and cosine decay."""
        class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
            def __init__(
                self,
                total_steps: int,
                peak_lr: float,
                warmup_steps: int
            ):
                super().__init__()
                self.total_steps = tf.cast(total_steps, tf.float32)
                self.peak_lr = tf.cast(peak_lr, tf.float32)
                
                # Adjust warmup_steps to not exceed half of total_steps
                adjusted_warmup_steps = min(warmup_steps, max(1, total_steps // 10))
                self.warmup_steps = tf.cast(adjusted_warmup_steps, tf.float32)
                
                # Calculate and store constants
                self.initial_lr = self.peak_lr * 0.1  # Start at 10% of peak
                self.min_lr = self.peak_lr * 0.01     # Minimum 1% of peak
                
                logger.info(f"Learning rate schedule initialized:")
                logger.info(f"  Initial LR: {float(self.initial_lr):.6f}")
                logger.info(f"  Peak LR: {float(self.peak_lr):.6f}")
                logger.info(f"  Min LR: {float(self.min_lr):.6f}")
                logger.info(f"  Warmup steps: {int(self.warmup_steps)}")
                logger.info(f"  Total steps: {int(self.total_steps)}")
            
            def __call__(self, step):
                step = tf.cast(step, tf.float32)
                
                # Warmup phase
                warmup_factor = tf.minimum(1.0, step / self.warmup_steps)
                warmup_lr = self.initial_lr + (self.peak_lr - self.initial_lr) * warmup_factor
                
                # Decay phase
                decay_steps = tf.maximum(1.0, self.total_steps - self.warmup_steps)
                decay_factor = (step - self.warmup_steps) / decay_steps
                decay_factor = tf.minimum(tf.maximum(0.0, decay_factor), 1.0)  # Clip to [0,1]
                
                cosine_decay = 0.5 * (1.0 + tf.cos(tf.constant(math.pi) * decay_factor))
                decay_lr = self.min_lr + (self.peak_lr - self.min_lr) * cosine_decay
                
                # Choose between warmup and decay
                final_lr = tf.where(step < self.warmup_steps, warmup_lr, decay_lr)
                
                # Ensure learning rate is valid
                final_lr = tf.maximum(self.min_lr, final_lr)
                final_lr = tf.where(tf.math.is_finite(final_lr), final_lr, self.min_lr)
                
                return final_lr
            
            def get_config(self):
                return {
                    "total_steps": self.total_steps,
                    "peak_lr": self.peak_lr,
                    "warmup_steps": self.warmup_steps,
                }
        
        return CustomSchedule(total_steps, peak_lr, warmup_steps)

    def _cosine_similarity(self, emb1: np.ndarray, emb2: np.ndarray) -> np.ndarray:
        """Compute cosine similarity between two numpy arrays."""
        normalized_emb1 = emb1 / np.linalg.norm(emb1, axis=1, keepdims=True)
        normalized_emb2 = emb2 / np.linalg.norm(emb2, axis=1, keepdims=True)
        return np.dot(normalized_emb1, normalized_emb2.T)

    def chat(
        self,
        query: str,
        conversation_history: Optional[List[Tuple[str, str]]] = None,
        quality_checker: Optional['ResponseQualityChecker'] = None,
        top_k: int = 5,
    ) -> Tuple[str, List[Tuple[str, float]], Dict[str, Any]]:
        """
        Example chat method that always uses cross-encoder re-ranking 
        if self.reranker is available.
        """
        @self.run_on_device
        def get_response(self_arg, query_arg):  # Add parameters that match decorator's expectations
            # 1) Build conversation context string
            conversation_str = self_arg._build_conversation_context(query_arg, conversation_history)
            
            # 2) Retrieve + cross-encoder re-rank
            results = self_arg.retrieve_responses_cross_encoder(
                query=conversation_str,
                top_k=top_k,
                reranker=self_arg.reranker,
                summarizer=self_arg.summarizer,
                summarize_threshold=512
            )

            # 3) Handle empty or confidence
            if not results:
                return (
                    "I'm sorry, but I couldn't find a relevant response.",
                    [],
                    {}
                )

            if quality_checker:
                metrics = quality_checker.check_response_quality(query_arg, results)
                if not metrics.get('is_confident', False):
                    return (
                        "I need more information to provide a good answer. Could you please clarify?",
                        results,
                        metrics
                    )
                return results[0][0], results, metrics
            
            return results[0][0], results, {}
        
        return get_response(self, query)

    def _build_conversation_context(
        self, 
        query: str, 
        conversation_history: Optional[List[Tuple[str, str]]]
    ) -> str:
        """Build conversation context with better memory management."""
        if not conversation_history:
            return f"{self.special_tokens['user']} {query}"
            
        conversation_parts = []
        for user_txt, assistant_txt in conversation_history:
            conversation_parts.extend([
                f"{self.special_tokens['user']} {user_txt}",
                f"{self.special_tokens['assistant']} {assistant_txt}"
            ])
            
        conversation_parts.append(f"{self.special_tokens['user']} {query}")
        return "\n".join(conversation_parts)