from train import create_synthetic_data
from utils.tokenizer import SimpleTokenizer
from utils.data_loader import TextDataset, TextDataLoader

if __name__ == "__main__":
    train_texts, _ = create_synthetic_data()
    tokenizer = SimpleTokenizer(vocab_size=10000, min_frequency=2)
    tokenizer.train(train_texts)
    dataset = TextDataset(train_texts, tokenizer, max_length=512)
    loader = TextDataLoader(dataset, batch_size=8, shuffle=True)

    batch = next(iter(loader))
    print(type(batch))
    if isinstance(batch, tuple):
        print("tuple length:", len(batch))
        for idx, item in enumerate(batch):
            print(f"item {idx} type:", type(item), "shape:", getattr(item, "shape", None))
    else:
        print("not tuple")