added changes in main.py file

main.py

@@ -7,8 +7,9 @@ from sentence_transformers import  util
 app = FastAPI()
 
 @app.get("/embeddings")
-def display_embedding(text:str):
-    embedding = models.get_embedding(text)
+def display_embedding(prediction : Prediction):
+    message = prediction.message
+    embedding = models.get_embedding(message)
     dimension = len(embedding)
     return {"Dimension" : {dimension : embedding.tolist()}}
 
@@ -16,7 +17,7 @@ def display_embedding(text:str):
 def display_prediction(prediction : Prediction):
     message = prediction.message
     embedding = models.get_embedding([message])
-    loaded_model = models.loaded_model
+    loaded_model = models.load_model('log_reg_model.pkl')
     result = loaded_model.predict(embedding).tolist()
     return {"Prediction": f"{message} is a {result}"}
 

models.py

@@ -12,7 +12,7 @@ def get_embedding(text):
 
 def train_model():
     sample_data_df = pd.read_excel("sms_process_data_main.xlsx")
-    sample_data_df.dropna(subset=['MessageText', 'label'], inplace=True)  # Drop rows with missing labels or texts
+    sample_data_df.dropna(subset=['MessageText', 'label'], inplace=True)
     input = sample_data_df['MessageText']
     label = sample_data_df['label']
     X_train, X_test, y_train, y_test = train_test_split(input, label, test_size=0.2, random_state=42)
@@ -32,6 +32,4 @@ def load_model(filename):
     with open(filename, 'rb') as model_file:
         loaded_model = pickle.load(model_file)
     print(f"Model loaded from {filename}")
-    return loaded_model
-
-loaded_model = load_model('log_reg_model.pkl')
+    return loaded_model