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)
    



# Example usage:
# sentence_a = "The cat sat on the mat."
# sentence_b = "A feline rested on the rug."
# similarity = calculate_cosine_similarity(sentence_a, sentence_b)

# if similarity is not None:
#     print(f"Cosine similarity between sentence_a and sentence_b: {similarity}")

# sentence_c = "This is a completely different sentence."
# similarity_ac = calculate_cosine_similarity(sentence_a, sentence_c)

# if similarity_ac is not None:
#     print(
#         f"Cosine similarity between sentence_a and sentence_c: {similarity_ac}")


# # Using a different model (you'll need to install it if you haven't already):
# similarity_different_model = calculate_cosine_similarity(
#     sentence_a, sentence_b, model_name="all-MiniLM-L6-v2")
# if similarity_different_model is not None:
#     print(f"Cosine similarity (different model): {similarity_different_model}")

# # Example of error handling if model name is wrong
# similarity_error = calculate_cosine_similarity(
#     sentence_a, sentence_b, model_name="wrong-model-name")

# if similarity_error is not None:
#     print(f"Cosine similarity (wrong model name): {similarity_error}")