test.py

from transformers import AutoTokenizer
from scripts.dataset import SquadDataset

from transformers import AutoTokenizer
from torch.utils.data import DataLoader, random_split
from scripts.model import SquadModel  

def test_small_chunking():
    print("\n===== RUNNING CHUNKING TEST =====")

    examples = [
        {
            "qas_id": "q1",
            "question": "Where do I live?",
            "context": "I live in Germany. Paris kappa kiakf fjkfn fnskdj"*10,
            "answer_text": "Germany",  # ✅ Must exist in context!
            "answer_start": 68,  # ✅ Position of "Germany" in context
            "answer_end": 91,    # ✅ 10 + len("Germany")
            "is_impossible": False
        }
    ]

    tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

    dataset = SquadDataset(examples, tokenizer, max_length=40, stride=10)
    test_dataloader = DataLoader(dataset, batch_size=32, num_workers=8, pin_memory=True)

    dataset.print_chunks("q1", tokenizer)

if __name__ == "__main__":
    test_small_chunking()