added function for cosine similarity

src/api/main.py

@@ -1,9 +1,10 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 import os
 import sys
 src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
 sys.path.append(src_directory)
 from modules import encoding_model
+from schemas.schemas import CosineSimilarity
 
 app = FastAPI()
 
@@ -15,15 +16,33 @@ def home():
 @app.get("/dimention")
 def display_dimention(message : str = "Hello World"):
     try:
-        dimention  = encoding_model.get_label(message)
-        return dimention
+        no_of_dimention  = encoding_model.get_prediction(message)[0]
+        dimentions = encoding_model.get_prediction(message)[1]
+        return {"message" : {"Prediction":{no_of_dimention:dimentions}}}
     except Exception as e:
         return f"Unable to fetch the data {e}"
     
 @app.get("/prediction")
 def display_prediction(message : str = "Give me a sms to predict"):
     try:
-        prediction = encoding_model.get_prediction(message)
+        prediction = encoding_model.get_prediction(message)[2]
         return {"message" : f"Given sms is a {prediction}"}
     except Exception as e:
-        return f"Unable to fetch the data {e}"
+        return f"Unable to fetch the data {e}"
+    
+@app.post("/cosine_similarity")
+def display_similarity(similarity: CosineSimilarity):
+    try:
+        if not similarity.message_1 or not similarity.message_2:
+            raise HTTPException(status_code=400, detail="Both messages must be non-empty strings.")
+
+        cosine_similarity = encoding_model.get_cosine_similarity(similarity.message_1, similarity.message_2)
+
+        return {
+            "message_1": similarity.message_1,
+            "message_2": similarity.message_2,
+            "cosine_similarity": cosine_similarity
+        }
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Unable to calculate cosine similarity: {str(e)}")

src/modules/encoding_model.py

@@ -1,4 +1,4 @@
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer,util
 from sklearn.model_selection import train_test_split
 from sklearn.linear_model import LogisticRegression
 import pandas as pd
@@ -28,24 +28,20 @@ def train_model():
         logreg_model = LogisticRegression(max_iter=100)
         logreg_model.fit(X_train_embeddings, y_train)
 
-def get_label(message):
+def get_prediction(message):
     if logreg_model is None:
         raise ValueError("Model has not been trained yet. Please call train_model first.")
 
     new_embeddings = encoding_model.encode([message])
     array = np.array(new_embeddings)[0].tolist()
 
-    prediction = logreg_model.predict(new_embeddings).tolist()
-
     no_of_dimensions = len(new_embeddings[0])
     dimension_df = pd.DataFrame(array, columns=["Dimension"])
+    prediction = logreg_model.predict(new_embeddings).tolist()
     
-    return {"Prediction_Dimension": {no_of_dimensions: dimension_df}}
-
-def get_prediction(message):
-    if logreg_model is None:
-        raise ValueError("Model has not been trained yet. Please call train_model first.")
+    return no_of_dimensions, dimension_df, prediction
 
-    new_embeddings = encoding_model.encode([message])
-    prediction = logreg_model.predict(new_embeddings).tolist()
-    return prediction
+def get_cosine_similarity(msg_1: str, msg_2: str):
+    embeddings = encoding_model.encode([msg_1, msg_2])
+    similarity = util.cos_sim(embeddings[0], embeddings[1]).item()
+    return round(similarity, 4)

src/schemas/schemas.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
 
 class CosineSimilarity(BaseModel):
-    text_1 : str
-    text_1 : str
+    message_1 : str
+    message_2 : str