import pytest
from transformers import AutoModelForCausalLM, AutoTokenizer
from swarms.models.mpt import MPT7B


def test_mpt7b_init():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )

    assert isinstance(mpt, MPT7B)
    assert mpt.model_name == "mosaicml/mpt-7b-storywriter"
    assert mpt.tokenizer_name == "EleutherAI/gpt-neox-20b"
    assert isinstance(mpt.tokenizer, AutoTokenizer)
    assert isinstance(mpt.model, AutoModelForCausalLM)
    assert mpt.max_tokens == 150


def test_mpt7b_run():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )
    output = mpt.run("generate", "Once upon a time in a land far, far away...")

    assert isinstance(output, str)
    assert output.startswith("Once upon a time in a land far, far away...")


def test_mpt7b_run_invalid_task():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )

    with pytest.raises(ValueError):
        mpt.run("invalid_task", "Once upon a time in a land far, far away...")


def test_mpt7b_generate():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )
    output = mpt.generate("Once upon a time in a land far, far away...")

    assert isinstance(output, str)
    assert output.startswith("Once upon a time in a land far, far away...")


def test_mpt7b_batch_generate():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )
    prompts = ["In the deep jungles,", "At the heart of the city,"]
    outputs = mpt.batch_generate(prompts, temperature=0.7)

    assert isinstance(outputs, list)
    assert len(outputs) == len(prompts)
    for output in outputs:
        assert isinstance(output, str)


def test_mpt7b_unfreeze_model():
    mpt = MPT7B(
        "mosaicml/mpt-7b-storywriter", "EleutherAI/gpt-neox-20b", max_tokens=150
    )
    mpt.unfreeze_model()

    for param in mpt.model.parameters():
        assert param.requires_grad