kiln_ai.adapters.fine_tune.finetune_registry

 1from typing import Type
 2
 3from kiln_ai.adapters.fine_tune.base_finetune import BaseFinetuneAdapter
 4from kiln_ai.adapters.fine_tune.fireworks_finetune import FireworksFinetune
 5from kiln_ai.adapters.fine_tune.openai_finetune import OpenAIFinetune
 6from kiln_ai.adapters.fine_tune.together_finetune import TogetherFinetune
 7from kiln_ai.adapters.fine_tune.vertex_finetune import VertexFinetune
 8from kiln_ai.adapters.ml_model_list import ModelProviderName
 9
10finetune_registry: dict[ModelProviderName, Type[BaseFinetuneAdapter]] = {
11    ModelProviderName.openai: OpenAIFinetune,
12    ModelProviderName.fireworks_ai: FireworksFinetune,
13    ModelProviderName.together_ai: TogetherFinetune,
14    ModelProviderName.vertex: VertexFinetune,
15}