Spaces:
Sleeping
Sleeping
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() | |