ekaterinatao's picture
Create config_data/config.py
2e205ce verified
from dataclasses import dataclass
@dataclass
class TgBot:
token: str
@dataclass
class DataConfig:
dataset: str
cls_vec: str
@dataclass
class ModelConfig:
bi_checkpoint: str
cross_checkpoint: str
device: str
@dataclass
class Config:
tg_bot: TgBot
data: DataConfig
model: ModelConfig
def load_config() -> Config:
return Config(
tg_bot=TgBot(
token='BOT_TOKEN'
),
data=DataConfig(
dataset='ekaterinatao/house_md_context3',
cls_vec='ekaterinatao/house_md_cls_embeds'
),
model=ModelConfig(
bi_checkpoint='ekaterinatao/house-md-bot-bert-bi-encoder',
cross_checkpoint='ekaterinatao/house-md-bot-bert-cross-encoder',
device='cpu'
)
)