import os import gc import numpy as np import faiss import tensorflow as tf import h5py import math from tqdm import tqdm import json from pathlib import Path from typing import Union, Optional, List, Tuple, Generator from transformers import AutoTokenizer from typing import List, Tuple, Generator from gpu_monitor import GPUMemoryMonitor from logger_config import config_logger logger = config_logger(__name__) class TFDataPipeline: def __init__( self, config, tokenizer, encoder, index_file_path: str, response_pool: List[str], max_length: int, query_embeddings_cache: dict, neg_samples: int = 3, index_type: str = 'IndexFlatIP', nlist: int = 100, max_retries: int = 3 ): #self.embedding_batch_size = embedding_batch_size self.config = config self.tokenizer = tokenizer self.encoder = encoder self.index_file_path = index_file_path self.response_pool = response_pool self.max_length = max_length self.neg_samples = neg_samples self.query_embeddings_cache = query_embeddings_cache # In-memory cache for embeddings self.index_type = index_type self.nlist = nlist self.embedding_batch_size = 16 if len(response_pool) < 100 else 64 self.search_batch_size = 16 if len(response_pool) < 100 else 64 self.max_batch_size = 16 if len(response_pool) < 100 else 64 self.memory_monitor = GPUMemoryMonitor() self.max_retries = max_retries if os.path.exists(index_file_path): logger.info(f"Loading existing FAISS index from {index_file_path}...") self.index = faiss.read_index(index_file_path) self.validate_faiss_index() logger.info("FAISS index loaded and validated successfully.") else: # Initialize FAISS index dimension = self.encoder.config.embedding_dim self.index = faiss.IndexFlatIP(dimension) logger.info(f"Initialized FAISS IndexFlatIP with dimension {dimension}.") if not self.index.is_trained: # Train the index if it's not trained. # TODO: Replace 'dimension' with embedding size dimension = self.query_embeddings_cache[next(iter(self.query_embeddings_cache))].shape[0] self.index.train(np.array(list(self.query_embeddings_cache.values())).astype(np.float32)) self.index.add(np.array(list(self.query_embeddings_cache.values())).astype(np.float32)) def validate_faiss_index(self): """Validates that the FAISS index has the correct dimensionality.""" expected_dim = self.encoder.config.embedding_dim if self.index.d != expected_dim: logger.error(f"FAISS index dimension {self.index.d} does not match encoder embedding dimension {expected_dim}.") raise ValueError("FAISS index dimensionality mismatch.") logger.info("FAISS index dimension validated successfully.") def save_embeddings_cache_hdf5(self, cache_file_path: str): """Save the embeddings cache to an HDF5 file.""" with h5py.File(cache_file_path, 'w') as hf: for query, emb in self.query_embeddings_cache.items(): hf.create_dataset(query, data=emb) logger.info(f"Embeddings cache saved to {cache_file_path}.") def load_embeddings_cache_hdf5(self, cache_file_path: str): """Load the embeddings cache from an HDF5 file.""" with h5py.File(cache_file_path, 'r') as hf: for query in hf.keys(): self.query_embeddings_cache[query] = hf[query][:] logger.info(f"Embeddings cache loaded from {cache_file_path}.") def save_faiss_index(self, index_file_path: str): faiss.write_index(self.index, index_file_path) logger.info(f"FAISS index saved to {index_file_path}") def load_faiss_index(self, index_file_path: str): self.index = faiss.read_index(index_file_path) logger.info(f"FAISS index loaded from {index_file_path}") def save_tokenizer(self, tokenizer_dir: str): self.tokenizer.save_pretrained(tokenizer_dir) logger.info(f"Tokenizer saved to {tokenizer_dir}") def load_tokenizer(self, tokenizer_dir: str): self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir) logger.info(f"Tokenizer loaded from {tokenizer_dir}") def estimate_total_pairs(self, dialogues: List[dict]) -> int: """Estimate total number of training pairs including hard negatives.""" base_pairs = sum( len([ 1 for i in range(len(d.get('turns', [])) - 1) if (d['turns'][i].get('speaker') == 'user' and d['turns'][i+1].get('speaker') == 'assistant') ]) for d in dialogues ) # Account for hard negatives return base_pairs * (1 + self.neg_samples) @staticmethod def load_json_training_data(data_path: Union[str, Path], debug_samples: Optional[int] = None) -> List[dict]: """ Load training data from a JSON file. Args: data_path (Union[str, Path]): Path to the JSON file containing dialogues. debug_samples (Optional[int]): Number of samples to load for debugging. Returns: List[dict]: List of dialogue dictionaries. """ logger.info(f"Loading training data from {data_path}...") data_path = Path(data_path) if not data_path.exists(): logger.error(f"Data file {data_path} does not exist.") return [] with open(data_path, 'r', encoding='utf-8') as f: dialogues = json.load(f) if debug_samples is not None: dialogues = dialogues[:debug_samples] logger.info(f"Debug mode: Limited to {debug_samples} dialogues") logger.info(f"Loaded {len(dialogues)} dialogues.") return dialogues def collect_responses(self, dialogues: List[dict]) -> List[str]: """Extract unique assistant responses from dialogues.""" response_set = set() for dialogue in tqdm(dialogues, desc="Processing Dialogues", unit="dialogue"): turns = dialogue.get('turns', []) for turn in turns: speaker = turn.get('speaker') text = turn.get('text', '').strip() if speaker == 'assistant' and text: # Ensure we don't exclude valid shorter responses if len(text) <= self.max_length: response_set.add(text) logger.info(f"Collected {len(response_set)} unique assistant responses from dialogues.") return list(response_set) def _extract_pairs_from_dialogue(self, dialogue: dict) -> List[Tuple[str, str]]: """Extract query-response pairs from a dialogue.""" pairs = [] turns = dialogue.get('turns', []) for i in range(len(turns) - 1): current_turn = turns[i] next_turn = turns[i+1] if (current_turn.get('speaker') == 'user' and next_turn.get('speaker') == 'assistant' and 'text' in current_turn and 'text' in next_turn): query = current_turn['text'].strip() positive = next_turn['text'].strip() pairs.append((query, positive)) return pairs def _compute_and_index_response_embeddings(self): """ Computes embeddings for the response pool and adds them to the FAISS index with progress bars. """ logger.info("Computing embeddings for the response pool...") # Ensure all responses are strings if not all(isinstance(response, str) for response in self.response_pool): logger.error("All elements in response_pool must be strings.") raise ValueError("Invalid data type in response_pool.") # Tokenization logger.info("Tokenizing responses...") encoded_responses = self.tokenizer( self.response_pool, padding=True, truncation=True, max_length=self.max_length, return_tensors='tf' ) response_ids = encoded_responses['input_ids'] # Compute embeddings in batches with progress bar batch_size = getattr(self, 'embedding_batch_size', 64) # Default to 64 if not set total_responses = len(response_ids) logger.info(f"Computing embeddings in batches of {batch_size}...") embeddings = [] with tqdm(total=total_responses, desc="Computing Embeddings", unit="response") as pbar: for i in range(0, total_responses, batch_size): batch_ids = response_ids[i:i + batch_size] # Compute embeddings batch_embeddings = self.encoder(batch_ids, training=False).numpy() # Normalize embeddings for cosine similarity faiss.normalize_L2(batch_embeddings) embeddings.append(batch_embeddings) pbar.update(len(batch_ids)) if embeddings: embeddings = np.vstack(embeddings).astype(np.float32) # Add embeddings to FAISS index with progress bar logger.info(f"Adding {len(embeddings)} response embeddings to FAISS index...") # Determine number of batches for indexing index_batch_size = getattr(self, 'index_batch_size', 1000) # Adjust as needed total_embeddings = len(embeddings) num_index_batches = math.ceil(total_embeddings / index_batch_size) with tqdm(total=total_embeddings, desc="Indexing Embeddings", unit="embedding") as pbar_index: for i in range(0, total_embeddings, index_batch_size): batch = embeddings[i:i + index_batch_size] self.index.add(batch) pbar_index.update(len(batch)) logger.info("Response embeddings added to FAISS index.") else: logger.warning("No embeddings to add to FAISS index.") # **Sanity Check:** Verify the number of embeddings in FAISS index logger.info(f"Total embeddings in FAISS index after addition: {self.index.ntotal}") # def _compute_and_index_response_embeddings(self): # """ # Computes embeddings for the response pool and adds them to the FAISS index. # """ # logger.info("Computing embeddings for the response pool...") # # Ensure all responses are strings # if not all(isinstance(response, str) for response in self.response_pool): # logger.error("All elements in response_pool must be strings.") # raise ValueError("Invalid data type in response_pool.") # # Proceed with tokenization # encoded_responses = self.tokenizer( # self.response_pool, # padding=True, # truncation=True, # max_length=self.max_length, # return_tensors='tf' # ) # response_ids = encoded_responses['input_ids'] # # Compute embeddings in batches # batch_size = getattr(self, 'embedding_batch_size', 64) # Default to 64 if not set # embeddings = [] # for i in range(0, len(response_ids), batch_size): # batch_ids = response_ids[i:i+batch_size] # # Compute embeddings # batch_embeddings = self.encoder(batch_ids, training=False).numpy() # # Normalize embeddings if using inner product or cosine similarity # faiss.normalize_L2(batch_embeddings) # embeddings.append(batch_embeddings) # if embeddings: # embeddings = np.vstack(embeddings).astype(np.float32) # # Add embeddings to FAISS index # logger.info(f"Adding {len(embeddings)} response embeddings to FAISS index...") # self.index.add(embeddings) # logger.info("Response embeddings added to FAISS index.") # else: # logger.warning("No embeddings to add to FAISS index.") # # **Sanity Check:** Verify the number of embeddings in FAISS index # logger.info(f"Total embeddings in FAISS index after addition: {self.index.ntotal}") def _find_hard_negatives_batch(self, queries: List[str], positives: List[str]) -> List[List[str]]: """Find hard negatives for a batch of queries with error handling and retries.""" retry_count = 0 total_responses = len(self.response_pool) # Set k to be neg_samples + additional candidates to improve negative selection k = self.neg_samples + 0 while retry_count < self.max_retries: try: # Compute embeddings in sub-batches to manage memory batch_size = 128 # Example sub-batch size; adjust as needed query_embeddings = [] for i in range(0, len(queries), batch_size): sub_queries = queries[i:i + batch_size] sub_embeddings = np.vstack([ self.query_embeddings_cache[q] for q in sub_queries ]).astype(np.float32) faiss.normalize_L2(sub_embeddings) query_embeddings.append(sub_embeddings) query_embeddings = np.vstack(query_embeddings) # Ensure contiguous memory layout query_embeddings = np.ascontiguousarray(query_embeddings) # Perform FAISS search on CPU distances, indices = self.index.search(query_embeddings, k) all_negatives = [] for query_indices, query, positive in zip(indices, queries, positives): negatives = [] positive_strip = positive.strip() seen = {positive_strip} for idx in query_indices: if idx >= 0 and idx < total_responses: candidate = self.response_pool[idx].strip() if candidate and candidate not in seen: seen.add(candidate) negatives.append(candidate) if len(negatives) >= self.neg_samples: break # If not enough negatives are found, pad with a special token while len(negatives) < self.neg_samples: negatives.append("") # Use a special token all_negatives.append(negatives) return all_negatives except KeyError as ke: retry_count += 1 logger.warning(f"Hard negative search attempt {retry_count} failed due to missing embeddings: {ke}") if retry_count == self.max_retries: logger.error("Max retries reached for hard negative search due to missing embeddings.") return [[""] * self.neg_samples for _ in queries] # Perform memory cleanup gc.collect() if tf.config.list_physical_devices('GPU'): tf.keras.backend.clear_session() except Exception as e: retry_count += 1 logger.warning(f"Hard negative search attempt {retry_count} failed: {e}") if retry_count == self.max_retries: logger.error("Max retries reached for hard negative search.") return [[""] * self.neg_samples for _ in queries] # Perform memory cleanup gc.collect() if tf.config.list_physical_devices('GPU'): tf.keras.backend.clear_session() def _tokenize_and_encode(self, queries: List[str], positives: List[str], negatives: List[List[str]]) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Tokenize and encode the queries, positives, and negatives. Returns: query_ids: [batch_size, max_length] positive_ids: [batch_size, max_length] negative_ids: [batch_size, neg_samples, max_length] """ # Tokenize queries q_enc = self.tokenizer( queries, padding="max_length", truncation=True, max_length=self.max_length, return_tensors="np" ) # Tokenize positives p_enc = self.tokenizer( positives, padding="max_length", truncation=True, max_length=self.max_length, return_tensors="np" ) # Tokenize negatives # Flatten negatives flattened_negatives = [neg for sublist in negatives for neg in sublist] if len(flattened_negatives) == 0: # No negatives at all: return a zero array n_ids = np.zeros((len(queries), self.neg_samples, self.max_length), dtype=np.int32) else: n_enc = self.tokenizer( flattened_negatives, padding="max_length", truncation=True, max_length=self.max_length, return_tensors="np" ) n_input_ids = n_enc["input_ids"] # Reshape to [batch_size, neg_samples, max_length] batch_size = len(queries) n_ids = n_input_ids.reshape(batch_size, self.neg_samples, self.max_length) # Convert to int32 query_ids = q_enc["input_ids"].astype(np.int32) positive_ids = p_enc["input_ids"].astype(np.int32) negative_ids = n_ids.astype(np.int32) return query_ids, positive_ids, negative_ids # Testing updated batch tokenization def prepare_and_save_data(self, dialogues: List[dict], tf_record_path: str, batch_size: int = 32): """ Processes dialogues in batches and saves to a TFRecord file using optimized batch tokenization and encoding. Args: dialogues (List[dict]): List of dialogue dictionaries. tf_record_path (str): Path to save the TFRecord file. batch_size (int): Number of dialogues to process per batch. """ logger.info(f"Preparing and saving data to {tf_record_path}...") num_dialogues = len(dialogues) num_batches = math.ceil(num_dialogues / batch_size) with tf.io.TFRecordWriter(tf_record_path) as writer: # Initialize progress bar with tqdm(total=num_batches, desc="Preparing Data Batches", unit="batch") as pbar: for i in range(num_batches): start_idx = i * batch_size end_idx = min(start_idx + batch_size, num_dialogues) batch_dialogues = dialogues[start_idx:end_idx] # Extract all query-positive pairs in the batch queries = [] positives = [] for dialogue in batch_dialogues: pairs = self._extract_pairs_from_dialogue(dialogue) for query, positive in pairs: if len(query) <= self.max_length and len(positive) <= self.max_length: queries.append(query) positives.append(positive) if not queries: pbar.update(1) continue # Skip if no valid queries # Compute and cache query embeddings try: self._compute_embeddings(queries) except Exception as e: logger.error(f"Error computing embeddings: {e}") pbar.update(1) continue # Skip to the next batch # Find hard negatives for the batch try: hard_negatives = self._find_hard_negatives_batch(queries, positives) except Exception as e: logger.error(f"Error finding hard negatives: {e}") pbar.update(1) continue # Skip to the next batch # Tokenize and encode all queries, positives, and negatives in the batch try: encoded_queries = self.tokenizer.batch_encode_plus( queries, max_length=self.config.max_context_token_limit, truncation=True, padding='max_length', return_tensors='tf' ) encoded_positives = self.tokenizer.batch_encode_plus( positives, max_length=self.config.max_context_token_limit, truncation=True, padding='max_length', return_tensors='tf' ) except Exception as e: logger.error(f"Error during tokenization: {e}") pbar.update(1) continue # Skip to the next batch # Flatten hard_negatives while maintaining alignment # Assuming hard_negatives is a list of lists, where each sublist corresponds to a query try: flattened_negatives = [neg for sublist in hard_negatives for neg in sublist] encoded_negatives = self.tokenizer.batch_encode_plus( flattened_negatives, max_length=self.config.max_context_token_limit, truncation=True, padding='max_length', return_tensors='tf' ) # Reshape encoded_negatives['input_ids'] to [num_queries, num_negatives, max_length] num_negatives = self.config.neg_samples reshaped_negatives = encoded_negatives['input_ids'].numpy().reshape(-1, num_negatives, self.config.max_context_token_limit) except Exception as e: logger.error(f"Error during negatives tokenization: {e}") pbar.update(1) continue # Skip to the next batch # Serialize each example and write to TFRecord for j in range(len(queries)): try: q_id = encoded_queries['input_ids'][j].numpy() p_id = encoded_positives['input_ids'][j].numpy() n_id = reshaped_negatives[j] feature = { 'query_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=q_id)), 'positive_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=p_id)), 'negative_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=n_id.flatten())), } example = tf.train.Example(features=tf.train.Features(feature=feature)) writer.write(example.SerializeToString()) except Exception as e: logger.error(f"Error serializing example {j} in batch {i}: {e}") continue # Skip to the next example # Update progress bar pbar.update(1) logger.info(f"Data preparation complete. TFRecord saved.") # def prepare_and_save_data(self, dialogues: List[dict], tfrecord_file_path: str, batch_size: int = 32): # """Processes dialogues in batches and saves to a TFRecord file.""" # with tf.io.TFRecordWriter(tfrecord_file_path) as writer: # total_dialogues = len(dialogues) # logger.debug(f"Total dialogues to process: {total_dialogues}") # with tqdm(total=total_dialogues, desc="Processing Dialogues", unit="dialogue") as pbar: # for i in range(0, total_dialogues, batch_size): # batch_dialogues = dialogues[i:i+batch_size] # # Process each batch_dialogues # # Extract pairs, find negatives, tokenize, and serialize # # Example: # for dialogue in batch_dialogues: # pairs = self._extract_pairs_from_dialogue(dialogue) # queries = [] # positives = [] # for query, positive in pairs: # queries.append(query) # positives.append(positive) # if queries: # # **Compute and cache query embeddings before searching** # self._compute_embeddings(queries) # # Find hard negatives # hard_negatives = self._find_hard_negatives_batch(queries, positives) # # for idx, negatives in enumerate(hard_negatives[:5]): # Log first 5 examples # # logger.debug(f"Query: {queries[idx]}") # # logger.debug(f"Positive: {positives[idx]}") # # logger.debug(f"Hard Negatives: {negatives}") # # Tokenize and encode # query_ids, positive_ids, negative_ids = self._tokenize_and_encode(queries, positives, hard_negatives) # # Serialize each example and write to TFRecord # for q_id, p_id, n_id in zip(query_ids, positive_ids, negative_ids): # feature = { # 'query_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=q_id)), # 'positive_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=p_id)), # 'negative_ids': tf.train.Feature(int64_list=tf.train.Int64List(value=n_id.flatten())), # } # example = tf.train.Example(features=tf.train.Features(feature=feature)) # writer.write(example.SerializeToString()) # pbar.update(len(batch_dialogues)) # logger.info(f"Data preparation complete. TFRecord saved at {tfrecord_file_path}") def _tokenize_negatives_tf(self, negatives): """Tokenizes negatives using tf.py_function.""" # Handle the case where negatives is an empty tensor if tf.size(negatives) == 0: return tf.zeros([0, self.neg_samples, self.max_length], dtype=tf.int32) # Convert EagerTensor to a list of strings negatives_list = [] for neg_list in negatives.numpy(): decoded_negs = [neg.decode("utf-8") for neg in neg_list if neg] # Filter out empty strings negatives_list.append(decoded_negs) # Flatten the list of lists flattened_negatives = [neg for sublist in negatives_list for neg in sublist] # Tokenize the flattened negatives if flattened_negatives: n_tokens = self.tokenizer( flattened_negatives, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='tf' ) # Reshape the tokens n_tokens_reshaped = tf.reshape(n_tokens['input_ids'], [-1, self.neg_samples, self.max_length]) return n_tokens_reshaped else: return tf.zeros([0, self.neg_samples, self.max_length], dtype=tf.int32) def _compute_embeddings(self, queries: List[str]) -> None: new_queries = [q for q in queries if q not in self.query_embeddings_cache] if not new_queries: return # All queries already cached # Compute embeddings for new queries new_embeddings = [] for i in range(0, len(new_queries), self.embedding_batch_size): batch_queries = new_queries[i:i + self.embedding_batch_size] encoded = self.tokenizer( batch_queries, padding=True, truncation=True, max_length=self.max_length, return_tensors='tf' ) batch_embeddings = self.encoder(encoded['input_ids'], training=False).numpy() faiss.normalize_L2(batch_embeddings) new_embeddings.extend(batch_embeddings) # Update the cache for query, emb in zip(new_queries, new_embeddings): self.query_embeddings_cache[query] = emb def data_generator(self, dialogues: List[dict]) -> Generator[Tuple[str, str, List[str]], None, None]: """ Generates training examples: (query, positive, hard_negatives). Wrapped the outer loop with tqdm for progress tracking. """ total_dialogues = len(dialogues) logger.debug(f"Total dialogues to process: {total_dialogues}") # Initialize tqdm progress bar with tqdm(total=total_dialogues, desc="Processing Dialogues", unit="dialogue") as pbar: for dialogue in dialogues: pairs = self._extract_pairs_from_dialogue(dialogue) for query, positive in pairs: # Ensure embeddings are computed, find hard negatives, etc. self._compute_embeddings([query]) hard_negatives = self._find_hard_negatives_batch([query], [positive])[0] yield (query, positive, hard_negatives) pbar.update(1) def _prepare_batch(self, queries: tf.Tensor, positives: tf.Tensor, negatives: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]: """Prepares a batch of data for training.""" # Convert EagerTensors to lists of strings queries_list = [query.decode("utf-8") for query in queries.numpy()] positives_list = [pos.decode("utf-8") for pos in positives.numpy()] # Tokenize queries and positives q_tokens = self.tokenizer(queries_list, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='tf') p_tokens = self.tokenizer(positives_list, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='tf') # Decode negatives and ensure they are lists of strings negatives_list = [] for neg_list in negatives.numpy(): decoded_negs = [neg.decode("utf-8") for neg in neg_list if neg] # Filter out empty strings negatives_list.append(decoded_negs) # Flatten negatives for tokenization if there are any valid negatives flattened_negatives = [neg for sublist in negatives_list for neg in sublist if neg] # Tokenize negatives if there are any n_tokens_reshaped = None if flattened_negatives: n_tokens = self.tokenizer(flattened_negatives, padding='max_length', truncation=True, max_length=self.max_length, return_tensors='tf') # Reshape n_tokens to match the expected shape based on the number of negatives per query # This part may need adjustment if the number of negatives varies per query n_tokens_reshaped = tf.reshape(n_tokens['input_ids'], [len(queries_list), -1, self.max_length]) else: # Create a placeholder tensor for the case where there are no negatives n_tokens_reshaped = tf.zeros([len(queries_list), 0, self.max_length], dtype=tf.int32) # Ensure n_tokens_reshaped has a consistent shape even when there are no negatives # Adjust shape to [batch_size, num_neg_samples, max_length] if n_tokens_reshaped.shape[1] != self.neg_samples: # Pad or truncate the second dimension to match neg_samples padding = tf.zeros([len(queries_list), tf.maximum(0, self.neg_samples - n_tokens_reshaped.shape[1]), self.max_length], dtype=tf.int32) n_tokens_reshaped = tf.concat([n_tokens_reshaped, padding], axis=1) n_tokens_reshaped = n_tokens_reshaped[:, :self.neg_samples, :] # Concatenate the positive and negative examples along the 'neg_samples' dimension combined_p_n_tokens = tf.concat([tf.expand_dims(p_tokens['input_ids'], axis=1), n_tokens_reshaped], axis=1) return q_tokens['input_ids'], combined_p_n_tokens def get_tf_dataset(self, dialogues: List[dict], batch_size: int) -> tf.data.Dataset: """ Creates a tf.data.Dataset for streaming training that yields (input_ids_query, input_ids_positive, input_ids_negatives). """ # 1) Start with a generator dataset dataset = tf.data.Dataset.from_generator( lambda: self.data_generator(dialogues), output_signature=( tf.TensorSpec(shape=(), dtype=tf.string), # Query (single string) tf.TensorSpec(shape=(), dtype=tf.string), # Positive (single string) tf.TensorSpec(shape=(self.neg_samples,), dtype=tf.string) # Hard Negatives (list of strings) ) ) # 2) Batch the raw strings dataset = dataset.batch(batch_size, drop_remainder=True) # 3) Map them through a tokenize step using `tf.py_function` dataset = dataset.map( lambda q, p, n: self._tokenize_triple(q, p, n), num_parallel_calls=1 #tf.data.AUTOTUNE ) dataset = dataset.prefetch(tf.data.AUTOTUNE) return dataset # def get_tf_dataset(self, dialogues: List[dict], batch_size: int) -> tf.data.Dataset: # """ # Creates a tf.data.Dataset for streaming training that yields # (input_ids_query, input_ids_positive, input_ids_negatives). # """ # # 1) Start with a generator dataset # dataset = tf.data.Dataset.from_generator( # lambda: self.data_generator(dialogues), # output_signature=( # tf.TensorSpec(shape=(), dtype=tf.string), # Query (single string) # tf.TensorSpec(shape=(), dtype=tf.string), # Positive (single string) # tf.TensorSpec(shape=(None,), dtype=tf.string) # Hard Negatives (list of strings) # ) # ) # # 2) Batch the raw strings # dataset = dataset.batch(batch_size) # # 3) Now map them through a tokenize step (via py_function) # dataset = dataset.map( # lambda q, p, n: self._tokenize_triple(q, p, n), # num_parallel_calls=1 #tf.data.AUTOTUNE # ) # dataset = dataset.prefetch(tf.data.AUTOTUNE) # return dataset def _tokenize_triple( self, q: tf.Tensor, p: tf.Tensor, n: tf.Tensor ) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]: """ Wraps a Python function via tf.py_function to convert tf.Tensors of strings -> Python lists of strings -> HF tokenizer -> Tensors of IDs. q is shape [batch_size], p is shape [batch_size], n is shape [batch_size, neg_samples] (i.e., each row is a list of negatives). """ # Use tf.py_function with limited parallelism q_ids, p_ids, n_ids = tf.py_function( func=self._tokenize_triple_py, inp=[q, p, n, tf.constant(self.max_length), tf.constant(self.neg_samples)], Tout=[tf.int32, tf.int32, tf.int32] ) # Manually set shape information q_ids.set_shape([None, self.max_length]) # [batch_size, max_length] p_ids.set_shape([None, self.max_length]) # [batch_size, max_length] n_ids.set_shape([None, self.neg_samples, self.max_length]) # [batch_size, neg_samples, max_length] return q_ids, p_ids, n_ids def _tokenize_triple_py( self, q: tf.Tensor, p: tf.Tensor, n: tf.Tensor, max_len: tf.Tensor, neg_samples: tf.Tensor ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """ Python function that: - Decodes each tf.string Tensor to a Python list of strings - Calls the HF tokenizer - Reshapes negatives - Returns np.array of int32s for (q_ids, p_ids, n_ids). q: shape [batch_size], p: shape [batch_size] n: shape [batch_size, neg_samples] max_len: scalar int neg_samples: scalar int """ max_len = int(max_len.numpy()) # Convert to Python int neg_samples = int(neg_samples.numpy()) # 1) Convert Tensors -> Python lists of strings q_list = [q_i.decode("utf-8") for q_i in q.numpy()] # shape [batch_size] p_list = [p_i.decode("utf-8") for p_i in p.numpy()] # shape [batch_size] # shape [batch_size, neg_samples], decode each row n_list = [] for row in n.numpy(): # row is shape [neg_samples], each is a tf.string decoded = [neg.decode("utf-8") for neg in row] n_list.append(decoded) # 2) Tokenize queries & positives q_enc = self.tokenizer( q_list, padding="max_length", truncation=True, max_length=max_len, return_tensors="np" ) p_enc = self.tokenizer( p_list, padding="max_length", truncation=True, max_length=max_len, return_tensors="np" ) # 3) Tokenize negatives # Flatten [batch_size, neg_samples] -> single list flattened_negatives = [neg for row in n_list for neg in row] if len(flattened_negatives) == 0: # No negatives at all: return a zero array n_ids = np.zeros((len(q_list), neg_samples, max_len), dtype=np.int32) else: n_enc = self.tokenizer( flattened_negatives, padding="max_length", truncation=True, max_length=max_len, return_tensors="np" ) # shape [batch_size * neg_samples, max_len] n_input_ids = n_enc["input_ids"] # We want to reshape to [batch_size, neg_samples, max_len] # Handle cases where there might be fewer negatives batch_size = len(q_list) n_ids_list = [] for i in range(batch_size): start_idx = i * neg_samples end_idx = start_idx + neg_samples row_negs = n_input_ids[start_idx:end_idx] # If fewer negatives, pad with zeros if row_negs.shape[0] < neg_samples: deficit = neg_samples - row_negs.shape[0] pad_arr = np.zeros((deficit, max_len), dtype=np.int32) row_negs = np.concatenate([row_negs, pad_arr], axis=0) n_ids_list.append(row_negs) # stack them -> shape [batch_size, neg_samples, max_len] n_ids = np.stack(n_ids_list, axis=0) # 4) Return as np.int32 arrays q_ids = q_enc["input_ids"].astype(np.int32) # shape [batch_size, max_len] p_ids = p_enc["input_ids"].astype(np.int32) # shape [batch_size, max_len] n_ids = n_ids.astype(np.int32) # shape [batch_size, neg_samples, max_len] return q_ids, p_ids, n_ids # 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 # def _find_hard_negatives_batch(self, queries: List[str], positives: List[str]) -> List[List[str]]: # """Find hard negatives for a batch of queries with error handling and retries.""" # retry_count = 0 # total_responses = len(self.response_pool) # while retry_count < self.max_retries: # try: # query_embeddings = np.vstack([ # self.query_embeddings_cache[q] for q in queries # ]).astype(np.float32) # query_embeddings = np.ascontiguousarray(query_embeddings) # faiss.normalize_L2(query_embeddings) # k = 1 # TODO: try higher k for better results # #logger.debug(f"Searching with k={k} among {total_responses} responses") # distances, indices = self.index.search(query_embeddings, k) # all_negatives = [] # for query_indices, query, positive in zip(indices, queries, positives): # negatives = [] # positive_strip = positive.strip() # seen = {positive_strip} # for idx in query_indices: # if idx >= 0 and idx < total_responses: # candidate = self.response_pool[idx].strip() # if candidate and candidate not in seen: # seen.add(candidate) # negatives.append(candidate) # if len(negatives) >= self.neg_samples: # break # # Pad with a special empty negative if necessary # while len(negatives) < self.neg_samples: # negatives.append("") # Use a special token # all_negatives.append(negatives) # return all_negatives