import gradio as gr
import numpy as np
import pickle
import joblib
import tensorflow as tf
from tensorflow.keras.preprocessing.sequence import pad_sequences

# === Load tokenizer ===
with open("tokenizer.pkl", "rb") as f:
    tokenizer = pickle.load(f)

# === Load label encoder ===
label_encoder = joblib.load("label_encoder.pkl")

# === Load TFLite model ===
interpreter = tf.lite.Interpreter(model_path="Model.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# === Predict function ===
def predict(text):
    # Preprocess input text
    sequence = tokenizer.texts_to_sequences([text])
    padded = pad_sequences(sequence, maxlen=input_details[0]['shape'][1])
    input_data = np.array(padded, dtype=np.float32)

    # Set input tensor
    interpreter.set_tensor(input_details[0]['index'], input_data)
    interpreter.invoke()

    # Get output tensor
    output = interpreter.get_tensor(output_details[0]['index'])[0]
    predicted_index = np.argmax(output)

    # Decode label
    predicted_label = label_encoder.inverse_transform([predicted_index])[0]
    return predicted_label

# === Gradio Interface ===
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="Enter text"),
    outputs=gr.Textbox(label="Predicted Label"),
    title="TFLite Text Classifier",
    description="Enter a sentence to classify using a TensorFlow Lite model."
)

iface.launch()