File size: 1,154 Bytes
a6bbf63
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
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
import uvicorn

app = FastAPI()

class PredictionRequest(BaseModel):
    sentence: str

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

@app.get("/")
def health_check():
    return {"Message": "Service is healthy", "Status": "OK"}

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

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