main.py

from fastapi import FastAPI
from pydantic import BaseModel
from components.model_loader import ModelLoader
from components.pipeline_preparer import PipelinePreparer
from components.predictor import Predictor
from utils.commons import setup_logging
import uvicorn

logger = setup_logging("main.log")

app = FastAPI()

class PredictionRequest(BaseModel):
    sentence: str

@app.on_event("startup")
async def startup_event():
    logger.info("Initializing model...")
    try:
        # Model initialization
        loader = ModelLoader()
        tokenizer, model = loader.load_model()
        pipeline = PipelinePreparer.prepare_pipeline(tokenizer, model)
        app.state.predictor = Predictor(pipeline)
        logger.info("Model initialized successfully")
    except Exception as e:
        logger.error(f"Error initializing model: {e}")
        app.state.predictor = None

@app.get("/")
def health_check():
    logger.info("Health check endpoint called")
    return {"Message": "Service is healthy", "Status": "OK"}

@app.post("/predict")
def predict(request: PredictionRequest):
    logger.info(f"Prediction request received: {request.sentence}")
    if not app.state.predictor:
        logger.error("Model not initialized")
        return {"error": "Model not initialized"}
    return app.state.predictor.predict(request.sentence)

if __name__ == "__main__":
    logger.info("Starting FastAPI app...")
    uvicorn.run(app, host="0.0.0.0", port=7860)