import os import numpy as np from transformers import TFAutoModel, AutoTokenizer import tensorflow as tf 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 import re from tf_data_pipeline import TFDataPipeline from response_quality_checker import ResponseQualityChecker from cross_encoder_reranker import CrossEncoderReranker from conversation_summarizer import DeviceAwareModel, Summarizer 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 = 10 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 and freeze layers based on config self.pretrained = TFAutoModel.from_pretrained(config.pretrained_model) self._freeze_layers() # Add Pooling layer (Global Average Pooling), Projection layer, Dropout, and Normalization self.pooler = tf.keras.layers.GlobalAveragePooling1D() self.projection = tf.keras.layers.Dense( config.embedding_dim, activation='tanh', name="projection", dtype=tf.float32 ) 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) 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, device: str = None, strategy=None, reranker: Optional[CrossEncoderReranker] = None, summarizer: Optional[Summarizer] = None, mode: str = 'training' ): super().__init__() self.config = config self.strategy = strategy self.device = device or self._setup_default_device() self.mode = mode.lower() # Initialize reranker, summarizer, tokenizer, encoder, and memory monitor self.reranker = reranker or self._initialize_reranker() self.tokenizer = self._initialize_tokenizer() self.encoder = self._initialize_encoder() self.summarizer = summarizer or self._initialize_summarizer() # Initialize data pipeline logger.info("Initializing TFDataPipeline.") self.data_pipeline = TFDataPipeline( config=self.config, tokenizer=self.tokenizer, encoder=self.encoder, index_file_path='new_iteration/data_prep_iterative_models/faiss_indices/faiss_index_production.index', response_pool=[], max_length=self.config.max_context_token_limit, query_embeddings_cache={}, neg_samples=self.config.neg_samples, index_type='IndexFlatIP', nlist=100, # Not used with IndexFlatIP max_retries=self.config.max_retries ) # Collect unique responses from dialogues if self.mode == 'inference': logger.info("Mode set to 'inference'. Loading FAISS index and response pool.") self._load_faiss_index_and_responses() elif self.mode != 'training': logger.error(f"Unsupported mode in RetrievalChatbot init: {self.mode}") raise ValueError(f"Unsupported mode in RetrievalChatbot init: {self.mode}") # Initialize training history self.history = { "train_loss": [], "val_loss": [], "train_metrics": {}, "val_metrics": {} } 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.""" return Summarizer( tokenizer=self.tokenizer, model_name="t5-small", max_summary_length=self.config.max_context_token_limit // 4, device=self.device, max_summary_rounds=2 ) 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": "", "assistant": "", "context": "", "sep": "" } tokenizer.add_special_tokens( {'additional_special_tokens': list(special_tokens.values())} ) return tokenizer def _initialize_encoder(self) -> EncoderModel: """Initialize the EncoderModel and resize token embeddings.""" logger.info("Initializing encoder model...") encoder = EncoderModel( self.config, name="shared_encoder", ) new_vocab_size = len(self.tokenizer) encoder.pretrained.resize_token_embeddings(new_vocab_size) logger.info(f"Token embeddings resized to: {new_vocab_size}") return encoder def _load_faiss_index_and_responses(self) -> None: """Load FAISS index and response pool for inference.""" try: logger.info(f"Loading FAISS index from {self.data_pipeline.index_file_path}...") self.data_pipeline.load_faiss_index(self.data_pipeline.index_file_path) logger.info("FAISS index loaded successfully.") # Load response pool associated with the FAISS index response_pool_path = self.data_pipeline.index_file_path.replace('.index', '_responses.json') if os.path.exists(response_pool_path): with open(response_pool_path, 'r', encoding='utf-8') as f: self.data_pipeline.response_pool = json.load(f) logger.info(f"Loaded {len(self.data_pipeline.response_pool)} responses from {response_pool_path}.") else: logger.error(f"Response pool file not found at {response_pool_path}.") raise FileNotFoundError(f"Response pool file not found at {response_pool_path}.") # Validate FAISS index and response pool self.data_pipeline.validate_faiss_index() logger.info("FAISS index and response pool validated successfully.") except Exception as e: logger.error(f"Failed to load FAISS index and response pool: {e}") raise @classmethod def load_model(cls, load_dir: Union[str, Path], mode: str = 'training') -> 'RetrievalChatbot': """ Load saved models and configuration. """ load_dir = Path(load_dir) # 1) Load config with open(load_dir / "config.json", "r") as f: config = ChatbotConfig.from_dict(json.load(f)) # 2) Initialize chatbot chatbot = cls(config, mode=mode) # 3) Load DistilBERT from huggingface folder chatbot.encoder.pretrained = TFAutoModel.from_pretrained( load_dir / "shared_encoder", config=config ) dummy_input = tf.zeros((1, config.max_context_token_limit), dtype=tf.int32) _ = chatbot.encoder(dummy_input, training=False) # 4) Load tokenizer chatbot.tokenizer = AutoTokenizer.from_pretrained(load_dir / "tokenizer") logger.info(f"Models and tokenizer loaded from {load_dir}") # 5) Load the custom top layers' weights custom_weights_path = load_dir / "encoder_custom_weights.weights.h5" if custom_weights_path.exists(): chatbot.encoder.load_weights(str(custom_weights_path)) logger.info("Loaded custom encoder weights for projection/dropout/etc.") else: logger.warning(f"No custom encoder weights found at {custom_weights_path}. The top-level projection layer won't have learned parameters.") # 6) If in inference mode, load FAISS, etc. if mode == 'inference': cls._prepare_model_for_inference(chatbot, load_dir) return chatbot @classmethod def _prepare_model_for_inference(cls, chatbot: 'RetrievalChatbot', load_dir: Path) -> None: """Internal method to load inference components.""" try: # Load FAISS index faiss_path = load_dir / 'faiss_indices/faiss_index_production.index' if faiss_path.exists(): chatbot.index = faiss.read_index(str(faiss_path)) logger.info("FAISS index loaded successfully") else: raise FileNotFoundError(f"FAISS index not found at {faiss_path}") # Load response pool response_pool_path = load_dir / 'faiss_indices/faiss_index_production_responses.json' if response_pool_path.exists(): with open(response_pool_path, 'r') as f: chatbot.response_pool = json.load(f) logger.info(f"Loaded {len(chatbot.response_pool)} responses") else: raise FileNotFoundError(f"Response pool not found at {response_pool_path}") # Verify dimensions match if chatbot.index.d != chatbot.config.embedding_dim: raise ValueError( f"FAISS index dimension {chatbot.index.d} doesn't match " f"model dimension {chatbot.config.embedding_dim}" ) except Exception as e: logger.error(f"Error loading inference components: {e}") raise 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 the HF DistilBERT submodule: self.encoder.pretrained.save_pretrained(save_dir / "shared_encoder") # ALSO save custom top-level layers' weights self.encoder.save_weights(save_dir / "encoder_custom_weights.weights.h5") # Save tokenizer self.tokenizer.save_pretrained(save_dir / "tokenizer") logger.info(f"Models and tokenizer saved to {save_dir}.") def sigmoid(self, x: float) -> float: return 1 / (1 + np.exp(-x)) def retrieve_responses_cross_encoder( self, query: str, top_k: int = 10, reranker: Optional[CrossEncoderReranker] = None, summarizer: Optional[Summarizer] = None, summarize_threshold: int = 512 ) -> List[Tuple[str, float]]: """ Retrieve top-k responses with optional domain-based boosting and cross-encoder re-ranking. Args: query: The user's input text. top_k: Number of final results to return. reranker: CrossEncoderReranker for refined scoring, if available. summarizer: Summarizer for long queries, if desired. summarize_threshold: Summarize if query wordcount > threshold. Returns: List of (response_text, final_score). """ # 1) Optional query summarization if summarizer and len(query.split()) > summarize_threshold: logger.info(f"Query is long ({len(query.split())} words). Summarizing.") query = summarizer.summarize_text(query) logger.info(f"Summarized Query: {query}") detected_domain = self.detect_domain_from_query(query) #logger.debug(f"Detected domain '{detected_domain}' for query: {query}") # Retrieve initial candidates from FAISS initial_k = min(top_k * 10, len(self.data_pipeline.response_pool)) faiss_candidates = self.retrieve_responses_faiss(query, domain=detected_domain, top_k=initial_k) texts = [item[0] for item in faiss_candidates] # Re-rank these boosted candidates if not reranker: reranker = CrossEncoderReranker(model_name="cross-encoder/ms-marco-MiniLM-L-12-v2") ce_scores = reranker.rerank(query, texts, max_length=256) # Combine cross-encoder score with the base FAISS score (simple multiplicative approach) final_candidates = [] for (resp_text, faiss_score), ce_score in zip(faiss_candidates, ce_scores): # TODO: dial this in. ce_prob = self.sigmoid(ce_score) # ~ relevance in [0..1] faiss_norm = (faiss_score + 1)/2.0 combined_score = 0.9 * ce_prob + 0.1 * faiss_norm # alpha = 0.9 # print(f'CE SCORE: {ce_score} FAISS SCORE: {faiss_score}') # combined_score = alpha * ce_score + (1 - alpha) * faiss_score length_adjusted_score = self.length_adjust_score(resp_text, combined_score) #combined_score = ce_score * faiss_score #final_candidates.append((resp_text, combined_score)) final_candidates.append((resp_text, length_adjusted_score)) # Sort descending by combined score final_candidates.sort(key=lambda x: x[1], reverse=True) # Return top_k return final_candidates[:top_k] DOMAIN_KEYWORDS = { 'restaurant': ['restaurant', 'dining', 'food', 'dine', 'reservation', 'table', 'menu', 'cuisine', 'eat', 'place to eat', 'hungry', 'chef', 'dish', 'meal', 'brunch', 'bistro', 'buffet', 'catering', 'gourmet', 'fast food', 'fine dining', 'takeaway', 'delivery', 'restaurant booking'], 'movie': ['movie', 'cinema', 'film', 'ticket', 'showtime', 'showing', 'theater', 'flick', 'screening', 'film ticket', 'film show', 'blockbuster', 'premiere', 'trailer', 'director', 'actor', 'actress', 'plot', 'genre', 'screen', 'sequel', 'animation', 'documentary'], 'ride_share': ['ride', 'taxi', 'uber', 'lyft', 'car service', 'pickup', 'dropoff', 'driver', 'cab', 'hailing', 'rideshare', 'ride hailing', 'carpool', 'chauffeur', 'transit', 'transportation', 'hail ride'], 'coffee': ['coffee', 'café', 'cafe', 'starbucks', 'espresso', 'latte', 'mocha', 'americano', 'barista', 'brew', 'cappuccino', 'macchiato', 'iced coffee', 'cold brew', 'espresso machine', 'coffee shop', 'tea', 'chai', 'java', 'bean', 'roast', 'decaf'], 'pizza': ['pizza', 'delivery', 'order food', 'pepperoni', 'topping', 'pizzeria', 'slice', 'pie', 'margherita', 'deep dish', 'thin crust', 'cheese', 'oven', 'tossed', 'sauce', 'garlic bread', 'calzone'], 'auto': ['car', 'vehicle', 'repair', 'maintenance', 'mechanic', 'oil change', 'garage', 'auto shop', 'tire', 'check engine', 'battery', 'transmission', 'brake', 'engine diagnostics', 'carwash', 'detail', 'alignment', 'exhaust', 'spark plug', 'dashboard'], } def extract_keywords(self, query: str) -> List[str]: """ Return any domain keywords present in the query (lowercased). """ query_lower = query.lower() found = set() for domain, kw_list in self.DOMAIN_KEYWORDS.items(): for kw in kw_list: if kw in query_lower: found.add(kw) return list(found) def length_adjust_score(self, text: str, base_score: float) -> float: """ Penalize very short lines or numeric lines; mildly reward longer lines. Adjust carefully so you don't overshadow cross-encoder signals. """ words = text.split() wcount = len(words) # Penalty if under 3 words if wcount < 4: return base_score * 0.8 # Bonus for lines > 12 words if wcount > 12: extra = min(wcount - 12, 8) bonus = 0.0005 * extra base_score += bonus return base_score def detect_domain_from_query(self, query: str) -> str: """ Detect the domain of the query based on keywords. """ domain_patterns = { 'restaurant': r'\b(restaurant|restaurants?|dining|food|foods?|dine|reservation|reservations?|table|tables?|menu|menus?|cuisine|cuisines?|eat|eats?|place\s?to\s?eat|places\s?to\s?eat|hungry|chef|chefs?|dish|dishes?|meal|meals?|fork|forks?|knife|knives?|spoon|spoons?|brunch|bistro|buffet|buffets?|catering|caterings?|gourmet|fast\s?food|fine\s?dining|takeaway|takeaways?|delivery|deliveries|restaurant\s?booking)\b', 'movie': r'\b(movie|movies?|cinema|cinemas?|film|films?|ticket|tickets?|showtime|showtimes?|showing|showings?|theater|theaters?|flick|flicks?|screening|screenings?|film\s?ticket|film\s?tickets?|film\s?show|film\s?shows?|blockbuster|blockbusters?|premiere|premieres?|trailer|trailers?|director|directors?|actor|actors?|actress|actresses?|plot|plots?|genre|genres?|screen|screens?|sequel|sequels?|animation|animations?|documentary|documentaries)\b', 'ride_share': r'\b(ride|rides?|taxi|taxis?|uber|lyft|car\s?service|car\s?services?|pickup|pickups?|dropoff|dropoffs?|driver|drivers?|cab|cabs?|hailing|hailings?|rideshare|rideshares?|ride\s?hailing|ride\s?hailings?|carpool|carpools?|chauffeur|chauffeurs?|transit|transits?|transportation|transportations?|hail\s?ride|hail\s?rides?)\b', 'coffee': r'\b(coffee|coffees?|café|cafés?|cafe|cafes?|starbucks|espresso|espressos?|latte|lattes?|mocha|mochas?|americano|americanos?|barista|baristas?|brew|brews?|cappuccino|cappuccinos?|macchiato|macchiatos?|iced\s?coffee|iced\s?coffees?|cold\s?brew|cold\s?brews?|espresso\s?machine|espresso\s?machines?|coffee\s?shop|coffee\s?shops?|tea|teas?|chai|chais?|java|javas?|bean|beans?|roast|roasts?|decaf)\b', 'pizza': r'\b(pizza|pizzas?|delivery|deliveries|order\s?food|order\s?foods?|pepperoni|pepperonis?|topping|toppings?|pizzeria|pizzerias?|slice|slices?|pie|pies?|margherita|margheritas?|deep\s?dish|deep\s?dishes?|thin\s?crust|thin\s?crusts?|cheese|cheeses?|oven|ovens?|tossed|tosses?|sauce|sauces?|garlic\s?bread|garlic\s?breads?|calzone|calzones?)\b', 'auto': r'\b(car|cars?|vehicle|vehicles?|repair|repairs?|maintenance|maintenances?|mechanic|mechanics?|oil\s?change|oil\s?changes?|garage|garages?|auto\s?shop|auto\s?shops?|tire|tires?|check\s?engine|check\s?engines?|battery|batteries?|transmission|transmissions?|brake|brakes?|engine\s?diagnostics|engine\s?diagnostic|carwash|carwashes?|detail|details?|alignment|alignments?|exhaust|exhausts?|spark\s?plug|spark\s?plugs?|dashboard|dashboards?)\b', } # Check for matches for domain, pattern in domain_patterns.items(): if re.search(pattern, query.lower()): return domain return 'other' def is_numeric_response(self, text: str) -> bool: """ Return True if `text` is purely digits (and/or spaces), with optional punctuation like '.' at the end. """ pattern = r'^[\s]*[\d]+([\s.,\d]+)*[\s]*$' return bool(re.match(pattern, text.strip())) def retrieve_responses_faiss( self, query: str, domain: str = 'other', top_k: int = 5, boost_factor: float = 1.05 ) -> List[Tuple[str, float]]: """ Retrieve top-k responses from the FAISS index (IndexFlatIP) given a user query. Args: query (str): The user input text. domain (str, optional): The detected domain. Defaults to 'other'. top_k (int, optional): Number of top results to return. Defaults to 5. boost_factor (float, optional): Factor to boost scores for keyword matches. Defaults to 1.3. Returns: List[Tuple[str, float]]: List of (response_text, similarity) sorted by descending similarity. """ # Encode the query q_emb = self.data_pipeline.encode_query(query) q_emb_np = q_emb.reshape(1, -1).astype('float32') # Search the index distances, indices = self.data_pipeline.index.search(q_emb_np, top_k * 10) # IndexFlatIP: 'distances' are inner products (cosine similarities for normalized vectors) candidates = [] for rank, idx in enumerate(indices[0]): if idx < 0: continue response = self.data_pipeline.response_pool[idx] text = response.get('text', '').strip() cand_domain = response.get('domain', 'other') score = distances[0][rank] # Skip purely numeric or extremely short text (fewer than 3 words): words = text.split() if len(words) < 4: continue if self.is_numeric_response(text): continue candidates.append((text, cand_domain, score)) if not candidates: logger.warning("No valid candidates found after initial numeric/length filtering.") return [] # Sort candidates by score descending candidates.sort(key=lambda x: x[2], reverse=True) # Filter in-domain responses in_domain = [c for c in candidates if c[1] == domain] if not in_domain: logger.info(f"No in-domain responses found for '{domain}'. Using all candidates.") in_domain = candidates # Boost responses containing query keywords query_keywords = self.extract_keywords(query) boosted = [] for (resp_text, resp_domain, score) in in_domain: new_score = score # If the domain is known AND the response text # shares any query keywords, apply a small boost if query_keywords and any(kw in resp_text.lower() for kw in query_keywords): new_score *= boost_factor #logger.debug(f"Boosting response: '{resp_text}' by factor {boost_factor}") # Apply length penalty/bonus new_score = self.length_adjust_score(resp_text, new_score) boosted.append((resp_text, new_score)) # Sort boosted responses boosted.sort(key=lambda x: x[1], reverse=True) # Print top 10 for resp, score in boosted[:150]: logger.debug(f"Candidate: '{resp}' with score {score}") # 8) Return top_k return boosted[:top_k] def chat( self, query: str, conversation_history: Optional[List[Tuple[str, str]]] = None, quality_checker: Optional['ResponseQualityChecker'] = None, top_k: int = 10, ) -> 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): # 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.tokenizer.additional_special_tokens[self.tokenizer.additional_special_tokens.index('')]} {query}" conversation_parts = [] for user_txt, assistant_txt in conversation_history: conversation_parts.extend([ f"{self.tokenizer.additional_special_tokens[self.tokenizer.additional_special_tokens.index('')]} {user_txt}", f"{self.tokenizer.additional_special_tokens[self.tokenizer.additional_special_tokens.index('')]} {assistant_txt}" ]) conversation_parts.append(f"{self.tokenizer.additional_special_tokens[self.tokenizer.additional_special_tokens.index('')]} {query}") return "\n".join(conversation_parts) def train_model( 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 = 1e-5, warmup_steps_ratio: float = 0.1, early_stopping_patience: int = 3, min_delta: float = 1e-4, test_mode: bool = False, initial_epoch: int = 0 ) -> None: """ Train the retrieval model using a pre-prepared TFRecord dataset. This method handles: - Checkpoint loading/restoring - LR scheduling - Epoch/iteration tracking - Optional training-history logging - Basic early stopping """ logger.info("Starting training with pre-prepared TFRecord dataset...") def parse_tfrecord_fn(example_proto, max_length, neg_samples): """ Parses a single TFRecord example. """ 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 # Count total records in 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 buffer_size = max(1, total_pairs // 10) # 10% of the dataset 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 & LR 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=tf.cast(peak_lr, tf.float32), 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=tf.cast(peak_lr, tf.float32)) logger.info("Using fixed learning rate.") # Initialize optimizer with dummy step dummy_input = tf.zeros((1, self.config.max_context_token_limit), dtype=tf.int32) with tf.GradientTape() as tape: dummy_output = self.encoder(dummy_input) dummy_loss = tf.cast(tf.reduce_mean(dummy_output), tf.float32) dummy_grads = tape.gradient(dummy_loss, self.encoder.trainable_variables) self.optimizer.apply_gradients(zip(dummy_grads, self.encoder.trainable_variables)) # Create checkpoint and manager checkpoint = tf.train.Checkpoint( epoch=tf.Variable(0, dtype=tf.int32), optimizer=self.optimizer, model=self.encoder ) manager = tf.train.CheckpointManager( checkpoint, directory=checkpoint_dir, max_to_keep=3, checkpoint_name='ckpt' ) # Restore from existing checkpoint if present latest_checkpoint = manager.latest_checkpoint history_path = Path(checkpoint_dir) / 'training_history.json' # If you want to log all epoch losses across runs if not hasattr(self, 'history'): self.history = {'train_loss': [], 'val_loss': [], 'learning_rate': []} if latest_checkpoint and not test_mode: # Add checkpoint inspection logger.info(f"\nTrying to load checkpoint from: {latest_checkpoint}") reader = tf.train.load_checkpoint(latest_checkpoint) # shape_from_key = reader.get_variable_to_shape_map() # dtype_from_key = reader.get_variable_to_dtype_map() # logger.info("\nCheckpoint Variables:") # for key in shape_from_key: # logger.info(f"{key}: dtype={dtype_from_key[key]} - Shape: {shape_from_key[key]}") status = checkpoint.restore(latest_checkpoint) status.assert_consumed() logger.info(f"Restored from checkpoint: {latest_checkpoint}") logger.info(f"Optimizer iterations after restore: {self.optimizer.iterations.numpy()}") # Verify learning rate after restore if use_lr_schedule: current_lr = float(lr_schedule(self.optimizer.iterations)) else: current_lr = float(self.optimizer.learning_rate.numpy()) logger.info(f"Current learning rate after restore: {current_lr:.2e}") # Derive initial_epoch from checkpoint name if not passed in ckpt_number = int(latest_checkpoint.split('ckpt-')[-1]) if initial_epoch == 0: initial_epoch = ckpt_number # Assign to checkpoint.epoch so we keep counting from that checkpoint.epoch.assign(tf.cast(initial_epoch, tf.int32)) logger.info(f"Resuming from epoch {initial_epoch}") # If you want to load old history from file: if history_path.exists(): try: with open(history_path, 'r') as f: self.history = json.load(f) logger.info(f"Loaded previous training history from {history_path}") except Exception as e: logger.warning(f"Could not load history, starting fresh: {e}") # Fix for custom weights not being saved in the full model. self.save_models(Path(checkpoint_dir) / "pretrained_full_model") logger.info(f"Manually saved custom weights after restore.") else: logger.info("Starting training from scratch") checkpoint.epoch.assign(tf.cast(0, tf.int32)) initial_epoch = 0 # Set up 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}") # Parse dataset dataset = tf.data.TFRecordDataset(tfrecord_file_path) # Optional: test/debug mode with small subset if test_mode: subset_size = 150 dataset = dataset.take(subset_size) logger.info(f"TEST MODE: Using only {subset_size} examples") # Recompute sizes, steps, epochs, etc., as needed total_pairs = subset_size train_size = int(total_pairs * (1 - validation_split)) val_size = total_pairs - train_size batch_size = min(batch_size, val_size) steps_per_epoch = math.ceil(train_size / batch_size) val_steps = math.ceil(val_size / batch_size) total_steps = steps_per_epoch * epochs buffer_size = max(1, total_pairs // 10) epochs = min(epochs, 5) # For quick debug early_stopping_patience = 2 logger.info(f"New training pairs: {train_size}") logger.info(f"New validation pairs: {val_size}") dataset = dataset.map( lambda x: parse_tfrecord_fn(x, self.config.max_context_token_limit, self.config.neg_samples), num_parallel_calls=tf.data.AUTOTUNE ) # Train/val split train_dataset = dataset.take(train_size) val_dataset = dataset.skip(train_size).take(val_size) # Shuffle and batch train_dataset = train_dataset.shuffle(buffer_size=buffer_size) train_dataset = train_dataset.batch(batch_size, drop_remainder=True) train_dataset = train_dataset.prefetch(tf.data.AUTOTUNE) val_dataset = val_dataset.batch(batch_size, drop_remainder=False) val_dataset = val_dataset.prefetch(tf.data.AUTOTUNE) val_dataset = val_dataset.cache() # Training loop best_val_loss = float("inf") epochs_no_improve = 0 for epoch in range(int(checkpoint.epoch.numpy()) + 1, epochs + 1): checkpoint.epoch.assign(epoch) logger.info(f"Starting Epoch {epoch}...") # --- Training Phase --- epoch_loss_avg = tf.keras.metrics.Mean(dtype=tf.float32) batches_processed = 0 # Progress bar 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 for q_batch, p_batch, n_batch in train_dataset: loss, grad_norm, post_clip_norm = 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(): step = (epoch - 1) * steps_per_epoch + batches_processed tf.summary.scalar("loss", tf.cast(loss, tf.float32), step=step) tf.summary.scalar("gradient_norm_pre_clip", tf.cast(grad_norm, tf.float32), step=step) tf.summary.scalar("gradient_norm_post_clip", tf.cast(post_clip_norm, tf.float32), step=step) # 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}", "pre_clip": f"{grad_norm.numpy():.2e}", "post_clip": f"{post_clip_norm.numpy():.2e}", "lr": f"{current_lr:.2e}", "batches": f"{batches_processed}/{steps_per_epoch}" }) gc.collect() # End the epoch early if we've processed all steps 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(dtype=tf.float32) 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 last_valid_val_loss = None valid_batches = False for q_batch, p_batch, n_batch in val_dataset: # If batch is too small, skip if tf.shape(q_batch)[0] < 2: logger.warning(f"Skipping validation batch of size {tf.shape(q_batch)[0]}") continue valid_batches = True val_loss = self.validation_step(q_batch, p_batch, n_batch) val_loss_avg(val_loss) last_valid_val_loss = 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}" }) gc.collect() if val_batches_processed >= val_steps: break if not valid_batches: # If no valid batch is found, fallback logger.warning("No valid validation batches in this epoch") if last_valid_val_loss is not None: val_loss = last_valid_val_loss val_loss_avg(val_loss) else: val_loss = epoch_loss_avg.result() val_loss_avg(val_loss) if is_tqdm_val and val_pbar: val_pbar.close() # End of epoch: final stats 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}") # TensorBoard epoch logs 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() # (Optional) Save model for quick testing/inference model_save_path = Path(checkpoint_dir) / f"model_epoch_{epoch}" self.save_models(model_save_path) logger.info(f"Saved model for epoch {epoch} at {model_save_path}") # Update local history self.history['train_loss'].append(train_loss) self.history['val_loss'].append(val_loss) self.history.setdefault('learning_rate', []).append(current_lr) def convert_to_py_floats(obj): if isinstance(obj, dict): return {k: convert_to_py_floats(v) for k, v in obj.items()} elif isinstance(obj, list): return [convert_to_py_floats(x) for x in obj] elif isinstance(obj, (np.float32, np.float64)): return float(obj) elif tf.is_tensor(obj): return float(obj.numpy()) else: return obj json_history = convert_to_py_floats(self.history) # Save training history to file every epoch # (Create or overwrite the file so we always have the latest.) with open(history_path, 'w') as f: json.dump(json_history, f) logger.info(f"Saved training history to {history_path}") # Early stopping 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.cast(tf.einsum('bd,bkd->bk', q_enc, combined_p_n), tf.float32) labels = tf.zeros([bs], dtype=tf.int32) # Keep labels as int32 loss = tf.nn.sparse_softmax_cross_entropy_with_logits( labels=labels, logits=dot_products ) loss = tf.cast(tf.reduce_mean(loss), tf.float32) # Calculate gradients gradients = tape.gradient(loss, self.encoder.trainable_variables) gradients_norm = tf.cast(tf.linalg.global_norm(gradients), tf.float32) max_grad_norm = tf.constant(1.5, dtype=tf.float32) gradients, _ = tf.clip_by_global_norm(gradients, max_grad_norm, gradients_norm) post_clip_norm = tf.cast(tf.linalg.global_norm(gradients), tf.float32) # Apply gradients self.optimizer.apply_gradients(zip(gradients, self.encoder.trainable_variables)) return loss, gradients_norm, post_clip_norm @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.cast(tf.einsum('bd,bkd->bk', q_enc, combined_p_n), tf.float32) labels = tf.zeros([bs], dtype=tf.int32) # Keep labels as int32 loss = tf.nn.sparse_softmax_cross_entropy_with_logits( labels=labels, logits=dot_products ) loss = tf.cast(tf.reduce_mean(loss), tf.float32) 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 = tf.cast(self.peak_lr * 0.1, tf.float32) self.min_lr = tf.cast(self.peak_lr * 0.01, tf.float32) 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.cast(tf.minimum(1.0, step / self.warmup_steps), tf.float32) warmup_lr = self.initial_lr + (self.peak_lr - self.initial_lr) * warmup_factor # Decay phase decay_steps = tf.cast(tf.maximum(1.0, self.total_steps - self.warmup_steps), tf.float32) decay_factor = tf.cast((step - self.warmup_steps) / decay_steps, tf.float32) decay_factor = tf.cast(tf.minimum(tf.maximum(0.0, decay_factor), 1.0), tf.float32) cosine_decay = tf.cast(0.5 * (1.0 + tf.cos(tf.constant(math.pi, dtype=tf.float32) * decay_factor)), tf.float32) 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)