Update app.py

app.py

@@ -3,28 +3,25 @@ import xgboost as xgb
 import tensorflow as tf
 from transformers import RobertaTokenizer, TFRobertaModel
 
-# Define model paths
-PU_MODEL_PATH = "trained_model_pu/pu_model.json"
+# Load PU model
+pu_tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
+roberta_model = TFRobertaModel.from_pretrained("roberta-base")
+roberta_model.trainable = False  # Freeze RoBERTa model
 
-# Load PU model from Hugging Face
-pu_tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
-roberta_model = TFRobertaModel.from_pretrained('roberta-base')
-roberta_model.trainable = False
-
-# Load XGBoost classifier
+# Load XGBoost classifier directly from the uploaded file
 pu_classifier = xgb.Booster()
-pu_classifier.load_model(PU_MODEL_PATH)
+pu_classifier.load_model("xgboost_spam_filter.model")  # Use the filename directly
 
-# PU model classification function
+# Classification function
 def classify_with_pu(text):
     inputs = pu_tokenizer(text, return_tensors="tf", truncation=True, max_length=128)
     embeddings = roberta_model(inputs).last_hidden_state[:, 0, :]
     dmatrix = xgb.DMatrix(embeddings.numpy())
     pu_probs = pu_classifier.predict(dmatrix)
-    return {"spam_probability": max(0, min(1, float(pu_probs[0]) if pu_probs.size > 0 else 0.5))}
+    return {"prediction": "Spam" if pu_probs[0] > 0.5 else "Not Spam", "probability": float(pu_probs[0])}
 
 # Create API
-iface = gr.Interface(fn=classify_with_pu, inputs=gr.Textbox(), outputs="json")
+iface = gr.Interface(fn=classify_with_pu, inputs=gr.Textbox(label="Enter Message"), outputs="json")
 
 # Launch API
 if __name__ == "__main__":