import gradio as gr
import xgboost as xgb
import tensorflow as tf
from transformers import RobertaTokenizer, TFRobertaModel

# 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 XGBoost classifier directly from the uploaded file
pu_classifier = xgb.Booster()
pu_classifier.load_model("xgboost_spam_filter.model")  # Use the filename directly

# 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 {"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(label="Enter Message"), outputs="json")

# Launch API
if __name__ == "__main__":
    iface.launch()