muryshev's picture
Update main.py
6666b3d verified
raw
history blame
2.48 kB
import logging
import os
from contextlib import asynccontextmanager # noqa: F401
from pathlib import Path
from typing import Annotated # noqa: F401
import dotenv
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from transformers import AutoModel, AutoTokenizer
from common import dependencies as DI # noqa: F401
from common.common import configure_logging
from common.configuration import Configuration
from routes.auth import router as auth_router
from routes.dataset import router as dataset_router
from routes.document import router as document_router
from routes.entity import router as entity_router
from routes.evaluation import router as evaluation_router
from routes.llm import router as llm_router
from routes.llm_config import router as llm_config_router
from routes.llm_prompt import router as llm_prompt_router
# Защита от автоудаления линтером
_ = DI
_ = Annotated
_ = asynccontextmanager
# Загружаем переменные из .env
dotenv.load_dotenv()
CONFIG_PATH = os.environ.get('CONFIG_PATH', 'config_dev.yaml')
print("config path: ")
print(CONFIG_PATH)
config = Configuration(CONFIG_PATH)
logger = logging.getLogger(__name__)
configure_logging(
level=config.common_config.log_level,
config_file_path=config.common_config.log_file_path,
)
# Костыль №47364: Удаляем файл статуса обработки датасета при запуске приложения
tmp_path = Path(os.environ.get("APP_TMP_PATH", '.')) / 'tmp.json'
tmp_path.unlink(missing_ok=True)
print("Downloading model to cache...")
AutoTokenizer.from_pretrained(config.db_config.search.vectorizer_path)
AutoModel.from_pretrained(config.db_config.search.vectorizer_path)
print("Model cached successfully.")
app = FastAPI(title="Assistant control panel")
origins = ["*"]
app.add_middleware(
CORSMiddleware,
allow_origins=origins,
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.include_router(llm_router)
app.include_router(dataset_router)
app.include_router(document_router)
app.include_router(llm_config_router)
app.include_router(llm_prompt_router)
app.include_router(entity_router)
app.include_router(evaluation_router)
app.include_router(auth_router)
if __name__ == "__main__":
uvicorn.run(
"main:app",
host="localhost",
port=8885,
reload=True,
workers=4
)