from inspect_ai.model import get_model
from inspect_ai.solver import generate, system_message
import asyncio
from bench_af._abstract.objects import ModelOrganism
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
import torch
hf_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct").to("cuda")
tokenizer = AutoTokenizer.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct"
)
print(len(hf_model.model.layers))
print(hf_model.config.num_hidden_layers)

# model = get_model(
#     model="hf/meta-llama/Llama-3.1-8B-Instruct", device="auto")

gen_config = GenerationConfig(
    return_dict_in_generate=True,
    output_hidden_states=True,
)


async def main() -> None:
    # outputs = await model.generate("hello")
    # print(outputs)
    model_inputs = tokenizer(
        ["The secret to baking a good cake is"], return_tensors="pt").to("cuda")

    outputs_hf = hf_model.generate(
        **model_inputs, max_length=30, generation_config=gen_config)
    # print(outputs_hf.hidden_states)
    print("==============GenerateDecoderOnlyOutput===========================")
    print(outputs_hf.sequences.shape)

    print(
        f"Token 1, layer 5 shape (batch_size, generated_length, hidden_size): {outputs_hf.hidden_states[1][5].shape}")
    print(
        f"Length of first token tuple (num layers): {len(outputs_hf.hidden_states[1])}")
    print(f" Num tokens generated {len(outputs_hf.hidden_states)}")
    print(f"Outputs model.generate : {type(outputs_hf)}")

    print("==============CausalLMOutptWithPast===========================")

    outputs_hf_forward = hf_model(
        **model_inputs, output_hidden_states=True,
        return_dict=True
    )
    print(
        f"Layer 5 with shape (batch_size, generated_length, hidden_size): {outputs_hf_forward.hidden_states[5].shape}")
    print(
        f"Length of first (and only) token tuple: {len(outputs_hf_forward.hidden_states)}")
    print(f"Outputs model.forward: {type(outputs_hf_forward)}")


asyncio.run(main())

# model_organism = ModelOrganism(
#     name="discriminatory",
#     model=model,
#     solver=[system_message(""), generate()],
#     supported_environments={"discriminatory", "mvp"},
# )


# def get() -> ModelOrganism:
#     return model_organism
