from ctransformers import AutoModelForCausalLM


class TestModel:
    def test_generate(self, lib):
        llm = AutoModelForCausalLM.from_pretrained("marella/gpt-2-ggml", lib=lib)
        response = llm("AI is going to", seed=5, max_new_tokens=3)
        assert response == " be a big"

        token = llm.sample()
        logits = llm.logits
        value = logits[token]
        logits[token] -= 1
        assert logits[token] == llm.logits[token] == value - 1
        llm.logits[token] *= 2
        assert logits[token] == llm.logits[token] == (value - 1) * 2

        assert llm.eos_token_id == 50256
        assert llm.vocab_size == 50257
        assert llm.context_length == 1024