Api/app/main.py

import logging
from contextlib import asynccontextmanager

from app.modules.clearVariables.routes.clearVariables_route import \
    router as clear_variables_routes
from app.modules.documentHandeler.routes.document_handeler_route import \
    router as upload_file_routes
from app.modules.querySearch.routes.querySearch_route import \
    router as query_search_routes
from fastapi import APIRouter, FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from transformers import (AutoModel, AutoModelForMaskedLM, AutoTokenizer,
                          pipeline)


@asynccontextmanager
async def lifespan(app: FastAPI):
    dense_model_name = "sentence-transformers/all-MiniLM-L6-v2"
    sparse_model_name = "prithivida/Splade_PP_en_v1"
    qa_model_name = "deepset/roberta-base-squad2"

    dense_tokenizer = AutoTokenizer.from_pretrained(dense_model_name)
    dense_model = AutoModel.from_pretrained(dense_model_name)

    sparse_tokenizer = AutoTokenizer.from_pretrained(sparse_model_name)
    sparse_model = AutoModelForMaskedLM.from_pretrained(sparse_model_name)

    qa_pipeline = pipeline("question-answering", model=qa_model_name)

    yield {
        "dense_tokenizer": dense_tokenizer,
        "dense_model": dense_model,
        "sparse_tokenizer": sparse_tokenizer,
        "sparse_model": sparse_model,
        "qa_pipeline": qa_pipeline,
    }


app = FastAPI(lifespan=lifespan)

origins = [
    "http://localhost:8000",
    "http://localhost:3000",
    "localhost:8000",
    "https://abadesalex-docurag.hf.space/api",
    "https://abadesalex-docurag.hf.space",
]


app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


app_router = APIRouter(prefix="/api")
app_router.include_router(upload_file_routes, prefix="/document", tags=["document"])
app_router.include_router(query_search_routes, prefix="/query", tags=["query"])
app_router.include_router(clear_variables_routes, prefix="/clear", tags=["clear"])


# @app_router.get("/")
# async def root():
#     return {"message": "Hello World"}


# Serve static files from the 'out/_next/static' directory
app.mount("/_next/static", StaticFiles(directory="app/out/_next/static"), name="static")


# Serve the main index.html
@app.get("/")
def read_root():
    return FileResponse("app/out/index.html")


@app.on_event("startup")
async def startup_event():
    logging.info("Application is starting up...")


@app.on_event("shutdown")
async def shutdown_event():
    logging.info("Application is shutting down...")


app.include_router(app_router)