import json
import pathlib

from .factory import build_ablated_model_class

_MODELS_JSON = pathlib.Path(__file__).parent / 'models.json'


def build_ablated_models() -> dict[str, type]:
  with open(_MODELS_JSON) as f:
    models: dict[str, str | dict] = json.load(f)

  ablated_models = {}

  for model_name, model_info in models.items():
    attention_classes = model_info['attention_classes']
    mlp_classes = model_info['mlp_classes']
    decoder_class = model_info['decoder_class']
    module_path = model_info['module_path']
    additional_subclasses = model_info.get('additional_subclasses', [])

    # Dynamically build the ablated model class and add it to globals so it can be imported by vLLM's ModelRegistry
    ablated_models[f"Ablated{model_name}"] = build_ablated_model_class(
        causal_lm_class=model_name,
        decoder_class=decoder_class,
        attention_classes=attention_classes,
        mlp_classes=mlp_classes,
        module_path=module_path,
        additional_subclasses=additional_subclasses,
    )

  print('Ablated models built:')
  print(ablated_models)

  return ablated_models


ablated_models = build_ablated_models()

globals().update(ablated_models)
__all__ = list(ablated_models.keys())
