from pathlib import Path

from inspect_ai.model import Model, get_model
from inspect_ai.solver import generate

from bench_af._abstract.objects import ModelOrganism
from nnsight import LanguageModel

CURRENT_DIRECTORY = Path(__file__).parent


def get() -> ModelOrganism:
    nnsight_model = LanguageModel("meta-llama/Llama-3.1-8B-Instruct", device_map="auto")

    return ModelOrganism(
        name=CURRENT_DIRECTORY.name,
        model=get_model(
            model="nnsight/test_model",
            nnsight_model=nnsight_model,
            nnsight_args={},
        ),
        solver=[generate()],
        supported_environments={"self_exfiltrate", "mvp"}
    )
