import gradio as gr
from sklearn.pipeline import Pipeline
import joblib


class CustomTextClassificationPipeline(Pipeline):
    def __init__(self):
        tfidf_vectorizer = joblib.load("tfidf_vectorizer.joblib")
        linear_svc = joblib.load("model_linear_svc.joblib")
        super().__init__([
            ('tfidf', tfidf_vectorizer),
            ('classifier', linear_svc)
        ])

    def predict(self, text):
        # Call the parent predict method to get the list of predicted labels
        y_pred_list = super().predict([text])

        # Convert the list to a string by taking the first element
        y_pred_str = str(y_pred_list[0])

        return y_pred_str


model = CustomTextClassificationPipeline()


def classify(sentence):
    return model.predict(sentence)


demo = gr.Interface(fn=classify, inputs="text", outputs="text")
demo.launch()