deploy

main.py

@@ -10,4 +10,4 @@ app.include_router(sms_router)
 @app.get("/")
 def home():
     
-    return {"message": "Welcome to embedding sms API, use /docs to post SMS text and get demenstions"}
+    return {"message": "Welcome to embedding sms API, use /docs to test endpoints"}

routes/sms_router.py

@@ -1,43 +1,47 @@
 from fastapi import APIRouter, HTTPException
-from schemas.schema import SMSRequest, EmbeddingResponse, SimilarityRequest, SimilarityResponse
+from schemas.schema import (
+    SMSRequest,
+    EmbeddingResponse,
+    SimilarityRequest,
+    SimilarityResponse,
+    PredictionRequest,
+    PredictionResponse
+)
 from service.embedded_service import generate_embeddings
-import numpy as np  # Import NumPy for cosine similarity calculation
+from service.prediction_service import predict_label
+import numpy as np
 
 # Initialize Router
 router = APIRouter()
 
 @router.post("/get_embeddings/", response_model=EmbeddingResponse)
 async def get_embeddings(sms_request: SMSRequest):
-    # Check if the input list is not empty
     if not sms_request.messages:
         raise HTTPException(status_code=400, detail="No messages provided.")
     
-    # Generate embeddings
     embeddings = generate_embeddings(sms_request.messages)
     
-    # Check if embeddings are generated and are in the correct format
     if not embeddings or not all(isinstance(emb, list) for emb in embeddings):
         raise HTTPException(status_code=500, detail="Failed to generate embeddings.")
     
-    # Get the dimensions from the first embedding (assuming all are the same)
     dimensions = len(embeddings[0]) if embeddings else 0
     
-    # Return the response as per the EmbeddingResponse model
     return EmbeddingResponse(dimensions=dimensions, embeddings=embeddings)
 
 @router.post("/calculate_similarity/", response_model=SimilarityResponse)
 async def calculate_similarity(similarity_request: SimilarityRequest):
-    # Get embeddings for both messages
     embeddings = generate_embeddings([similarity_request.message1, similarity_request.message2])
 
-    # Check if embeddings are generated for both messages
     if len(embeddings) != 2:
         raise HTTPException(status_code=500, detail="Failed to generate embeddings for both messages.")
 
-    # Calculate cosine similarity
     vec1 = np.array(embeddings[0])
     vec2 = np.array(embeddings[1])
     cosine_similarity = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
     
-    # Return response using the SimilarityResponse model
     return SimilarityResponse(similarity_score=float(cosine_similarity))
+
+@router.post("/predict_label/", response_model=PredictionResponse)
+async def predict_sms_label(prediction_request: PredictionRequest):
+    label, probability = predict_label(prediction_request.message)
+    return PredictionResponse(label=label, probability=probability)

schemas/schema.py

@@ -25,4 +25,11 @@ class SimilarityRequest(BaseModel):
     message2: str
 
 class SimilarityResponse(BaseModel):
-    similarity_score: float
+    similarity_score: float
+    
+class PredictionRequest(BaseModel):
+    message: str
+
+class PredictionResponse(BaseModel):
+    label: str
+    probability: float

service/embedded_service.py

@@ -5,4 +5,4 @@ def generate_embeddings(messages: list):
     # Generate embeddings
     embeddings = model.encode(messages)
     embeddings = np.array(embeddings).tolist()  # Convert to list for JSON serialization
-    return embeddings
+    return embeddings

service/prediction_service.py

@@ -0,0 +1,16 @@
+import pickle
+from sentence_transformers import SentenceTransformer
+import numpy as np
+
+# Load Model and Transformer
+with open('models/logistic_regression_model.pkl', 'rb') as f:
+    logistic_model = pickle.load(f)
+
+model = SentenceTransformer('models/sentence_transformer')
+
+def predict_label(message: str):
+    embedding = model.encode([message])
+    prediction = logistic_model.predict(embedding)[0]
+    probability = logistic_model.predict_proba(embedding)[0].max()
+    
+    return prediction, float(probability)

service/train_model.py

@@ -0,0 +1,25 @@
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, classification_report
+import pickle
+from sentence_transformers import SentenceTransformer
+
+file_name = "data/sms_process_data_main.xlsx"
+sheet = "Sheet1"
+df = pd.read_excel(file_name, sheet_name=sheet)
+
+X_train, X_test, y_train, y_test = train_test_split(df['MessageText'], df['label'], test_size=0.2, random_state=42)
+
+model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
+
+X_train_embeddings = model.encode(X_train.tolist())
+X_test_embeddings = model.encode(X_test.tolist())
+
+logistic_model = LogisticRegression(max_iter=100)
+logistic_model.fit(X_train_embeddings, y_train)
+
+with open('models/logistic_regression_model.pkl', 'wb') as f:
+    pickle.dump(logistic_model, f)
+
+model.save('models/sentence_transformer')