import yaml

from typing import Dict
from pydantic import BaseModel, ValidationError



class PileConfig(BaseModel):
    file2persona: Dict[str, str]
    file2prefix: Dict[str, str]
    persona2system: Dict[str, str]
    prompt: str

class InferenceConfig(BaseModel):
    chat_template: str

class RepoConfig(BaseModel):
    name: str
    tag: str

class ModelConfig(BaseModel):
    pile: PileConfig
    inference: InferenceConfig
    repo: RepoConfig

    @classmethod
    def from_yaml(cls, yaml_file = "datasets/config.yaml"):
        with open(yaml_file, 'r') as file:
            data = yaml.safe_load(file)
        try:
            return cls(**data)
        except ValidationError as e:
            raise e