senga-dnotes / config.py
serenarolloh's picture
Update config.py
0669c63 verified
raw
history blame
729 Bytes
from pydantic import BaseSettings
class Settings(BaseSettings):
dataset: str "senga-ml/dnotes-dataset-v6"
base_config: str = "naver-clova-ix/donut-base"
base_processor: str = "naver-clova-ix/donut-base"
base_model: str = "naver-clova-ix/donut-base"
def get_model_config(self, shipper_id: str):
if shipper_id in ["61", "81", "139"]:
model = "senga-ml/donut-16"
processor = "senga-ml/donut-16"
elif shipper_id in ["165", "127", "145"]:
model = "senga-ml/donut-17"
processor = "senga-ml/donut-17"
else:
model = self.base_model
processor = self.base_processor
return model, processor
settings = Settings()