scripts/test_dataloader.py

"""
Test script for MSC DataLoader in SHINE-LR-v3
Validates data processing pipeline and PyTorch Dataset integration
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data.msc_processor import MSCDataProcessor
from data.shine_dataset import SHINEDataset, shine_collate_fn
from torch.utils.data import DataLoader
from transformers import AutoTokenizer


def print_separator(char="=", length=60):
    """Print visual separator"""
    print(char * length)


def display_conversation(conv):
    """Display detailed information about a conversation"""
    print(f"Conv ID: {conv.conv_id}")
    print(f"Number of sessions: {len(conv.sessions)}")
    
    # Display init personas
    if conv.init_personas:
        print(f"Init personas: {conv.init_personas}")
    
    # Display each session
    for sess in conv.sessions:
        print(f"\n--- Session {sess.session_id} ---")
        
        # Truncate long strings for readability
        persona_preview = sess.persona_summary[:100] + "..." if len(sess.persona_summary) > 100 else sess.persona_summary
        print(f"Persona Summary: {persona_preview}")
        
        dialog_preview = sess.dialog_text[:150] + "..." if len(sess.dialog_text) > 150 else sess.dialog_text
        print(f"Dialog Preview: {dialog_preview}")
        
        if sess.followup:
            print(f"Followup: {sess.followup}")
        if sess.newfact:
            print(f"New Fact: {sess.newfact}")


def test_data_processing(msc_root: str, split: str = "valid"):
    """Test MSC data processing pipeline"""
    print_separator()
    print("🚀 Testing MSC DataLoader for SHINE-LR-v3")
    print_separator()
    
    if not os.path.exists(msc_root):
        print(f"❌ Directory not found: {msc_root}")
        return None
    
    # Load and process data
    processor = MSCDataProcessor(msc_root)
    conversations = processor.process_split(split=split)
    
    if len(conversations) == 0:
        print("❌ No conversations processed.")
        return None
    
    # Display sample conversation
    print_separator()
    print("📝 SAMPLE CONVERSATION")
    print_separator()
    display_conversation(conversations[0])
    
    return conversations


def test_pytorch_dataset(conversations, batch_size: int = 2):
    """Test PyTorch Dataset and DataLoader integration"""
    print_separator()
    print("🔄 Testing PyTorch Dataset")
    print_separator()
    
    # Initialize tokenizer
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Create dataset (use subset for testing)
    dataset = SHINEDataset(
        conversations=conversations[:10],
        tokenizer=tokenizer,
        max_context_len=512,
        max_qa_len=256
    )
    
    # Create dataloader
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=True,
        collate_fn=shine_collate_fn
    )
    
    # Test one batch
    for batch in dataloader:
        print(f"\nBatch size: {len(batch)}")
        
        for i, item in enumerate(batch):
            print(f"\n  Sample {i+1}:")
            print(f"    Conv ID: {item['conv_id']}")
            print(f"    Num sessions: {item['num_sessions']}")
            
            # Display session 1 info
            sess1 = item['sessions'][0]
            context_preview = sess1['context_text'][:100] + "..." if len(sess1['context_text']) > 100 else sess1['context_text']
            print(f"    Session 1 context preview: {context_preview}")
            print(f"    Session 1 QA input_ids shape: {sess1['qa']['input_ids'].shape}")
            
            persona_preview = sess1['target_persona'][:80] + "..." if len(sess1['target_persona']) > 80 else sess1['target_persona']
            print(f"    Session 1 target persona: {persona_preview}")
        
        break  # Only test first batch
    
    print_separator()
    print("✅ DataLoader test completed successfully!")
    print_separator()


def main():
    """Main test function"""
    msc_root = "/home/minhhieuano/2222/workspace/LearnAI/shine-lr/data/msc_raw_data/msc"
    
    # Test data processing
    conversations = test_data_processing(msc_root, split="valid")
    
    if conversations is None:
        return
    
    # Test PyTorch integration
    test_pytorch_dataset(conversations, batch_size=2)


if __name__ == "__main__":
    main()