from fastapi import FastAPI
import joblib
from sentence_transformers import SentenceTransformer

app = FastAPI()
model = SentenceTransformer(
    'Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)


@app.get("/")
def root():
    return {"message": "Welcom to FastAPI with Logistic Regression"}


@app.post("/dimensions/")
def get_dimension(message: str):
    message_embedding = model.encode([message])
    return {"dimensions": message_embedding.shape[1]}


@app.post("/classify/")
def get_dimension(message: str):
    loaded_model = joblib.load('sms_classifier_model.pkl')
    message_embedding = model.encode([message])
    prediction = loaded_model.predict(message_embedding)
    return {"Predicted Category": f"{prediction[0]}"}