Upload get_text_embeedings.py

get_text_embeedings.py

@@ -0,0 +1,113 @@
+"""
+Pre-computes T5 text embeddings for the MotionStreamer dataset.
+This script is GUARANTEED to be correct because it imports the
+original base Text2MotionDataset class and scans its internal 'data_dict'
+to discover every possible caption, including all sub-clips.
+"""
+import os
+import torch
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from tqdm import tqdm
+import argparse
+import hashlib
+import json
+import sys
+
+# --- CRITICAL: Import the actual dataset class from your project ---
+try:
+    from humanml3d_272.dataset_TM_train_motionstreamer import Text2MotionDataset
+except ImportError as e:
+    print("FATAL ERROR: Could not import the 'Text2MotionDataset' class.")
+    print("Please make sure you run this script from the root of your project directory.")
+    print(f"Original error: {e}")
+    sys.exit(1)
+
+def get_args_parser():
+    parser = argparse.ArgumentParser(description='Pre-compute T5 Embeddings')
+    
+    # --- Args needed for the Text2MotionDataset ---
+    # These must match the values your dataloader uses
+    parser.add_argument('--dataname', type=str, default='t2m_babel_272', help='Dataset name (for dataset init)')
+    parser.add_argument('--latent_dir', type=str, default='babel_272_stream/t2m_babel_latents', help='Latent dir (for dataset init)')
+    parser.add_argument('--unit_length', type=int, default=4, help='Unit length (for dataset init, 4 is a common default)')
+    
+    # --- Args for this script ---
+    parser.add_argument('--output_file', type=str, default='babel_272_stream/text_embeddings.npy', help='Path to save the output .npy file.')
+    parser.add_argument('--t5_model_path', type=str, default='sentence-t5-xl', help='Path or HF name for the Sentence-T5-XL model')
+    parser.add_argument('--batch_size', type=int, default=256, help='Batch size for T5 encoding.')
+    return parser
+
+def main():
+    parser = get_args_parser()
+    args = parser.parse_args()
+    print(f"Configuration:\n{json.dumps(vars(args), indent=4, sort_keys=True)}")
+
+    output_dir = os.path.dirname(args.output_file)
+    os.makedirs(output_dir, exist_ok=True)
+    print(f"Embeddings will be saved to: {args.output_file}")
+
+    # Load T5 model
+    print(f"Loading T5 model from: {args.t5_model_path}")
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    
+    if device == 'cuda' and torch.cuda.is_bf16_supported():
+        print("bfloat16 is supported, loading model in bf16.")
+        t5_model = SentenceTransformer(args.t5_model_path, device=device, model_kwargs={'torch_dtype': torch.bfloat16})
+    else:
+        print("bfloat16 not supported or not on CUDA, loading model in fp32.")
+        t5_model = SentenceTransformer(args.t5_model_path, device=device)
+    t5_model.eval()
+    for p in t5_model.parameters():
+        p.requires_grad = False
+    print("T5 model loaded successfully.")
+
+    # --- THIS IS THE CORRECT LOGIC (FROM YOUR OLD SCRIPT) ---
+    print("Instantiating the Text2MotionDataset to scan for all captions...")
+    dataset = Text2MotionDataset(
+        dataset_name=args.dataname, 
+        latent_dir=args.latent_dir,
+        unit_length=args.unit_length
+    )
+    
+    if not hasattr(dataset, 'data_dict') or not isinstance(dataset.data_dict, dict):
+        print("FATAL ERROR: The imported Text2MotionDataset does not have 'data_dict'.")
+        sys.exit(1)
+        
+    unique_captions = set()
+    unique_captions.add('') # Add the unconditional caption for 10% dropout
+    print("Extracting all unique captions from the dataset's internal dictionary...")
+    
+    for data_item in tqdm(dataset.data_dict.values(), desc="Scanning discovered samples"):
+        for text_dict in data_item['text']:
+            unique_captions.add(text_dict['caption'])
+    
+    captions_list = list(unique_captions)
+    print(f"Found {len(captions_list)} unique captions to encode.")
+    # --- END CORRECT LOGIC ---
+
+    print(f"Encoding {len(captions_list)} captions in batches of {args.batch_size}...")
+    with torch.no_grad():
+        all_embeddings = t5_model.encode(
+            captions_list,
+            batch_size=args.batch_size,
+            convert_to_tensor=True,
+            show_progress_bar=True
+        )
+
+    all_embeddings_fp32 = all_embeddings.to(torch.float32).cpu().numpy()
+
+    # Populate the dictionary
+    embeddings_dict = {}
+    for i, final_caption in enumerate(captions_list):
+        # Use sha256 to match the hash your *new* script was trying to use
+        caption_hash = hashlib.sha256(final_caption.encode('utf-8')).hexdigest()
+        embeddings_dict[caption_hash] = all_embeddings_fp32[i]
+
+    print(f"Saving {len(embeddings_dict)} embeddings to {args.output_file}...")
+    np.save(args.output_file, embeddings_dict, allow_pickle=True)
+
+    print("--- All text embeddings have been pre-computed and saved. ---")
+
+if __name__ == '__main__':
+    main()