from sentence_transformers import SentenceTransformer, util
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 classify(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]}"}


@app.post("/cosine-similarity")
def calculate_cosine_similarity(sentence1:str, sentence2:str):

    embeddings = model.encode([sentence1, sentence2])

    cosine_sim = util.cos_sim(embeddings[0], embeddings[1])
    # return str(cosine_sim[0])
    return round(cosine_sim.item(), 3)