modified code

Dockerfile

@@ -1,26 +1,18 @@
-# Use an official Python base image
-FROM python:3.10
+# Use official Python image as the base
+FROM python:3.9-slim
 
-# Set the working directory
+# Set working directory
 WORKDIR /app
 
-# Set environment variables for HF cache
-ENV HF_HOME="/app/cache"
-ENV TRANSFORMERS_CACHE="/app/cache"
-ENV SENTENCE_TRANSFORMERS_HOME="/app/cache"
-
-# Create the cache directory with appropriate permissions
-RUN mkdir -p /app/cache && chmod -R 777 /app/cache
-
-# Copy the requirements file and install dependencies
+# Copy requirements file and install dependencies
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Copy the application code
+# Copy the rest of the application
 COPY . .
 
-# Expose FastAPI default port
-EXPOSE 7860
+# Expose the FastAPI app port
+EXPOSE 8000
 
-# Run FastAPI with Uvicorn
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# Run the application with Uvicorn
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

main.py

@@ -1,32 +1,50 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import List, Tuple
-import numpy as np
-from sentence_transformers import SentenceTransformer
+from typing import List
+from model import get_embeddings, predict_sms_category
+from service import calculate_cosine_similarity
 
-# Load the pre-trained model
-model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
+# FastAPI app
+app = FastAPI()
 
-# Define request model
 class MessageRequest(BaseModel):
     messages: List[str]
 
-# Define response model
+class CosineSimilarityRequest(BaseModel):
+    message1: str
+    message2: str
+
+class PredictionRequest(BaseModel):
+    message: str
+
 class EmbeddingResponse(BaseModel):
-    dimensions: int  # Only return embedding size
+    dimensions: int
     numeric_values: List[List[float]]
 
-# Initialize FastAPI app
-app = FastAPI()
+class CosineSimilarityResponse(BaseModel):
+    similarity: float
+
+class PredictionResponse(BaseModel):
+    label: str
 
 @app.get("/")
-def home ():
-    return {"Message":"Welcome to homepage, kindly proceed by giving /docs in the URL" }
+def home():
+    return {"Message": "Welcome to the SMS classifier API. Use /docs for documentation."}
 
 @app.post("/embed", response_model=EmbeddingResponse)
 def embed(request: MessageRequest):
-    new_embeddings = model.encode(request.messages, convert_to_tensor=True)
+    embeddings = get_embeddings(request.messages)
     return EmbeddingResponse(
-        dimensions=new_embeddings.shape[1],  # Return only the embedding dimension
-        numeric_values=new_embeddings.tolist()
+        dimensions=embeddings.shape[1],  # Number of embedding dimensions
+        numeric_values=embeddings.tolist()
     )
+
+@app.post("/cosine_similarity", response_model=CosineSimilarityResponse)
+def cosine_similarity(request: CosineSimilarityRequest):
+    similarity = calculate_cosine_similarity(request.message1, request.message2)
+    return CosineSimilarityResponse(similarity=similarity)
+
+@app.post("/predict", response_model=PredictionResponse)
+def predict(request: PredictionRequest):
+    label = predict_sms_category(request.message)
+    return PredictionResponse(label=label)

model.py

@@ -0,0 +1,65 @@
+import joblib
+from sentence_transformers import SentenceTransformer
+import numpy as np
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import LabelEncoder
+from sklearn.model_selection import train_test_split
+import pandas as pd
+
+# Load pre-trained Sentence Transformer model
+model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
+
+# Load and preprocess SMS data (from an Excel file)
+def load_sms_data(file_path="data/sms_process_data_main.xlsx"):
+    data = pd.read_excel(file_path)
+    texts = data['MessageText'].tolist()
+    labels = data['label'].tolist()
+
+    embeddings = model.encode(texts, convert_to_tensor=True)
+    embeddings = embeddings.detach().numpy()
+
+    label_encoder = LabelEncoder()
+    encoded_labels = label_encoder.fit_transform(labels)
+
+    return embeddings, encoded_labels, label_encoder
+
+# Train and save the Logistic Regression model
+def train_sms_classifier():
+    embeddings, labels, label_encoder = load_sms_data()
+    X_train, X_test, y_train, y_test = train_test_split(embeddings, labels, test_size=0.3, random_state=42)
+
+    # Train Logistic Regression
+    lr_model = LogisticRegression()
+    lr_model.fit(X_train, y_train)
+
+    accuracy = lr_model.score(X_test, y_test)
+    print(f"Model Accuracy: {accuracy * 100:.2f}%")
+    
+    # Save the trained model and label encoder
+    joblib.dump(lr_model, 'model/sms_classifier_model.pkl')
+    joblib.dump(label_encoder, 'model/label_encoder.pkl')
+
+    return lr_model, label_encoder
+
+# Load the saved model and label encoder
+def load_saved_model():
+    lr_model = joblib.load('model/sms_classifier_model.pkl')
+    label_encoder = joblib.load('model/label_encoder.pkl')
+    return lr_model, label_encoder
+
+# Generate embeddings for the messages
+def get_embeddings(messages):
+    embeddings = model.encode(messages, convert_to_tensor=True)
+    return embeddings.detach().numpy()
+
+# Predict the label of an SMS message
+def predict_sms_category(message):
+    # Load the saved model and label encoder
+    lr_model, label_encoder = load_saved_model()
+
+    embedding = model.encode([message], convert_to_tensor=True)
+    embedding = embedding.detach().numpy()
+
+    prediction = lr_model.predict(embedding)
+    label = label_encoder.inverse_transform(prediction)[0]
+    return label

requirements.txt

@@ -1,6 +1,7 @@
 fastapi
 uvicorn
 pandas
-scikit-learn
 sentence-transformers
-numpy
+scikit-learn
+openpyxl
+joblib

service.py

@@ -0,0 +1,9 @@
+from sklearn.metrics.pairwise import cosine_similarity
+from model import model
+
+# Calculate cosine similarity between two messages
+def calculate_cosine_similarity(message1, message2):
+    embeddings = model.encode([message1, message2], convert_to_tensor=True)
+    embeddings = embeddings.detach().numpy()
+    similarity = cosine_similarity([embeddings[0]], [embeddings[1]])[0][0]
+    return similarity