config.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131
  1. import os
  2. from functools import lru_cache
  3. from pathlib import Path
  4. from pydantic import Field, field_validator
  5. from pydantic_settings import (
  6. BaseSettings,
  7. EnvSettingsSource,
  8. SettingsConfigDict,
  9. )
  10. class EnvSettingsSourceWithCommaLists(EnvSettingsSource):
  11. """Override decode_complex_value to handle comma-separated lists."""
  12. def decode_complex_value(self, field_name, field, value):
  13. if field_name == "backend_cors_origins" and isinstance(value, str):
  14. return [v.strip() for v in value.split(",") if v.strip()]
  15. return super().decode_complex_value(field_name, field, value)
  16. class Settings(BaseSettings):
  17. model_config = SettingsConfigDict(
  18. env_file=str(Path(__file__).resolve().parents[2] / ".env"),
  19. env_file_encoding="utf-8",
  20. case_sensitive=False,
  21. extra="ignore",
  22. )
  23. @classmethod
  24. def settings_customise_sources(
  25. cls, settings_cls, init_settings, env_settings, dotenv_settings, file_secret_settings
  26. ):
  27. return (
  28. init_settings,
  29. dotenv_settings,
  30. EnvSettingsSourceWithCommaLists(settings_cls),
  31. file_secret_settings,
  32. )
  33. # --- 数据路径 ---
  34. data_dir: Path = Path("/root/Fine-tuning/backend/data")
  35. # --- HuggingFace / ModelScope ---
  36. hf_token: str = ""
  37. hf_endpoint: str = "https://huggingface.co"
  38. use_modelscope: bool = False
  39. modelscope_endpoint: str = "https://modelscope.cn"
  40. # --- GPU / 硬件 ---
  41. cuda_visible_devices: str = "0"
  42. max_memory_per_gpu: str = "0"
  43. use_unsloth: bool = False
  44. # --- 后端 ---
  45. backend_host: str = "0.0.0.0"
  46. backend_port: int = 8000
  47. backend_env: str = "production"
  48. backend_log_level: str = "INFO"
  49. backend_cors_origins: list[str] = ["http://192.168.91.253:5173"]
  50. # --- 数据库 ---
  51. database_url: str = "sqlite+aiosqlite:///root/Fine-tuning/backend/data/finetuning.db"
  52. # --- 训练默认参数 ---
  53. default_peft_method: str = "lora"
  54. default_epochs: int = 3
  55. default_batch_size: int = 4
  56. default_gradient_accumulation: int = 4
  57. default_lr: float = 2e-4
  58. default_max_seq_length: int = 2048
  59. default_warmup_ratio: float = 0.05
  60. default_save_strategy: str = "epoch"
  61. default_eval_strategy: str = "epoch"
  62. default_eval_steps: int = 100
  63. # --- LoRA ---
  64. lora_r: int = 16
  65. lora_alpha: int = 32
  66. lora_dropout: float = 0.05
  67. lora_target_modules: str = "all-linear"
  68. # --- QLoRA ---
  69. qlora_bits: int = 4
  70. qlora_type: str = "nf4"
  71. qlora_double_quant: bool = True
  72. # --- 上传限制 ---
  73. max_upload_size_mb: int = 500
  74. allowed_dataset_formats: str = "jsonl,csv,parquet,json"
  75. @field_validator("backend_cors_origins", mode="before")
  76. @classmethod
  77. def parse_cors_origins(cls, v):
  78. if isinstance(v, str):
  79. return [origin.strip() for origin in v.split(",") if origin.strip()]
  80. return v
  81. @property
  82. def models_dir(self) -> Path:
  83. return self.data_dir / "models"
  84. @property
  85. def adapters_dir(self) -> Path:
  86. return self.data_dir / "adapters"
  87. @property
  88. def uploads_dir(self) -> Path:
  89. return self.data_dir / "uploads"
  90. @property
  91. def processed_dir(self) -> Path:
  92. return self.data_dir / "processed"
  93. def ensure_dirs(self) -> None:
  94. self.data_dir.mkdir(parents=True, exist_ok=True)
  95. for d in [self.models_dir, self.adapters_dir, self.uploads_dir, self.processed_dir]:
  96. d.mkdir(parents=True, exist_ok=True)
  97. @lru_cache
  98. def get_settings() -> Settings:
  99. settings = Settings()
  100. settings.ensure_dirs()
  101. # 设置 HF 环境变量
  102. if settings.hf_token:
  103. os.environ["HF_TOKEN"] = settings.hf_token
  104. os.environ["HF_ENDPOINT"] = settings.hf_endpoint
  105. if settings.cuda_visible_devices:
  106. os.environ["CUDA_VISIBLE_DEVICES"] = settings.cuda_visible_devices
  107. return settings