from fastapi import FastAPI
import models
import numpy as np
from schema import Prediction
from sentence_transformers import  util

app = FastAPI()

@app.get("/embeddings")
def display_embedding(text:str):
    embedding = models.get_embedding(text)
    dimension = len(embedding)
    return {"Dimension" : {dimension : embedding.tolist()}}

@app.post("/prediction")
def display_prediction(prediction : Prediction):
    message = prediction.message
    embedding = models.get_embedding([message])
    loaded_model = models.loaded_model
    result = loaded_model.predict(embedding).tolist()
    return {"Prediction": f"{message} is a {result}"}

@app.post("/cosine_similarity")
def display_cosine_similarity(prediction : Prediction):
    message = prediction.message
    message_1 = prediction.message_1
    embendding = models.get_embedding([message,message_1])
    similarity = util.cos_sim(embendding[0], embendding[1]).item()
    return {f"Cosine Similarity between {message}  and {message_1} is" : round(similarity, 4)}