borodache's picture
commit to change the repository to include already trained model and pipeline
1afb34d
raw
history blame contribute delete
867 Bytes
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()