Added ModelConfig

shared.py

@@ -0,0 +1,33 @@
+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