this repo has no description
hailey.at
1from typing import Dict, List
2from pydantic_settings import BaseSettings, SettingsConfigDict
3import torch
4
5
6class Config(BaseSettings):
7 data_dir: str = "./data"
8 nsfw_dir: str = "./dataset/nsfw"
9 sfw_humans_dir: str = "./dataset/sfw-human"
10 sfw_anime_dir: str = "./dataset/sfw-anime"
11
12 model_save_path: str = "./models/nsfw_vit_classifier"
13
14 model_name: str = "Falconsai/nsfw_image_detection"
15 num_classes: int = 2
16
17 batch_size: int = 16
18 num_epochs: int = 10
19 learning_rate: float = 5e-6
20 weight_decay: float = 0.01
21
22 train_ratio: float = 0.7
23 val_ratio: float = 0.15
24 test_ratio: float = 0.15
25
26 device: torch.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
27 num_workers: int = 4
28
29 class_names: List[str] = ["SFW", "NSFW"]
30 class_to_idx: Dict[str, int] = {"sfw": 0, "nsfw": 1}
31
32 firehose_unsure_output: str = "./firehose-unsure"
33 firehose_sure_output: str = "./firehose-sure"
34
35 kafka_bootstrap_server: str = "10.0.0.66:9092"
36 kafka_topic: str = "atproto-events"
37
38 ozone_pds_url: str = ""
39 ozone_did: str = ""
40 ozone_identifier: str = ""
41 ozone_password: str = ""
42
43 model_config = SettingsConfigDict(env_file=".env", env_file_encoding="utf-8")
44
45
46CONFIG = Config()