Update app.py

app.py

@@ -3,8 +3,6 @@ from pydantic import BaseModel
 import joblib
 import pandas as pd
 import re
-from sklearn.feature_extraction.text import TfidfVectorizer
-from sklearn.svm import LinearSVC
 from transformers import pipeline
 
 # Initialize FastAPI app
@@ -16,9 +14,8 @@ app = FastAPI(
     redoc_url="/redoc"
 )
 
-# Load model and vectorizer
+# Load the combined model pipeline (includes vectorizer)
 model = joblib.load("model.joblib")
-vectorizer = joblib.load("vectorizer.joblib")
 
 # Initialize NER pipeline
 ner = pipeline('ner', model='Davlan/xlm-roberta-base-ner-hrl', grouped_entities=True)
@@ -114,8 +111,7 @@ def restore_pii(masked_text, pii_map):
 def classify_email(data: EmailInput):
     raw_text = data.input_email_body
     masked_text, pii_map, entity_list = mask_and_store_all_pii(raw_text)
-    features = vectorizer.transform([masked_text])
-    predicted_category = model.predict(features)[0]
+    predicted_category = model.predict([masked_text])[0]
     return {
         "input_email_body": raw_text,
         "list_of_masked_entities": entity_list,
@@ -123,32 +119,6 @@ def classify_email(data: EmailInput):
         "category_of_the_email": predicted_category
     }
 
-# Retraining endpoint
-@app.post("/train")
-def train_model(new_example: TrainingExample):
-    df = pd.DataFrame([{"email_body": new_example.email_body, "label": new_example.label}])
-    try:
-        df.to_csv("training_data.csv", mode='a', header=not pd.io.common.file_exists("training_data.csv"), index=False)
-    except Exception as e:
-        return {"error": f"Failed to append to dataset: {str(e)}"}
-
-    # Load dataset
-    full_df = pd.read_csv("training_data.csv")
-    full_df['masked_text'] = full_df['email_body'].apply(lambda x: mask_and_store_all_pii(x)[0])
-
-    # Vectorize and train
-    new_vectorizer = TfidfVectorizer()
-    X = new_vectorizer.fit_transform(full_df['masked_text'])
-    y = full_df['label']
-    new_model = LinearSVC()
-    new_model.fit(X, y)
-
-    # Save updated model and vectorizer
-    joblib.dump(new_model, "model.joblib")
-    joblib.dump(new_vectorizer, "vectorizer.joblib")
-
-    return {"message": "Model retrained successfully with new example."}
-
 # Health check
 @app.get("/")
 def root():