updated

src/api/main.py

@@ -9,7 +9,7 @@ app = FastAPI()
 
 @app.get("/")
 def home():
-    encoding_model.train_model()
+    model = encoding_model.train_model()
     return {"message": "Welcome to Prediction Hub"}
 
 @app.get("/predict")

src/modules/encoding_model.py

@@ -1,5 +1,4 @@
 from sentence_transformers import SentenceTransformer
-model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
 from sklearn.model_selection import train_test_split
 from sklearn.linear_model import LogisticRegression
 import pandas as pd
@@ -8,6 +7,9 @@ import sys
 src_directory = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..", "src"))
 sys.path.append(src_directory)
 from data import sample_data
+import numpy as np
+
+model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
 
 encoding_model = model
 logreg_model = None
@@ -31,9 +33,10 @@ def get_label(message):
         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)
 
     no_of_dimensions = len(new_embeddings[0])
-    dimension_df = pd.DataFrame(new_embeddings[0], columns=["Dimension"])
-    return {"Prediction": prediction[0], "Prediction_Dimension": {no_of_dimensions: dimension_df}}
+    dimension_df = pd.DataFrame(array, columns=["Dimension"])
+    return {"Prediction_Dimension": {no_of_dimensions: dimension_df}}