"""Smoke tests for ai_toolkit.dl. Require torch and transformers (install with pip install -e ".[test]")."""

from ai_toolkit.dl import (
    BERTDataset,
    SelfAttention,
    SentimentLSTM,
    SentimentLSTMWithAttention,
    Vocabulary,
    get_bert_tokenizer,
)


def test_Vocabulary():
    vocab = Vocabulary().build([["a", "b", "a"], ["b", "c"]])
    assert len(vocab) >= 2
    ids = vocab.encode(["a", "b"])
    assert len(ids) == 2
    assert all(isinstance(i, int) for i in ids)


def test_SentimentLSTM():
    model = SentimentLSTM(vocab_size=100, num_classes=3)
    assert model is not None


def test_SelfAttention():
    attn = SelfAttention(64)
    assert attn is not None


def test_SentimentLSTMWithAttention():
    model = SentimentLSTMWithAttention(vocab_size=100)
    assert model is not None


def test_BERTDataset():
    ds = BERTDataset(["hello", "world"], max_length=8)
    assert len(ds) == 2
    item = ds[0]
    assert "input_ids" in item
    assert "attention_mask" in item


def test_get_bert_tokenizer():
    tok = get_bert_tokenizer("distilbert-base-uncased")
    assert tok is not None
