Update tasks/utils/predict.py

tasks/utils/predict.py

@@ -8,7 +8,7 @@ from tasks.utils.preprocessing import process_text
 import json
 from sklearn.feature_extraction.text import TfidfVectorizer
     
-def predict(input_df: pd.DataFrame, tfidf_path:str , model_path: str):
+def predict(input_df: pd.DataFrame, tfidf_path:str , tfidf_voc_path:str, model_path: str):
     """
     Predict the output using a saved TF-IDF vectorizer and Random Forest model.
 
@@ -31,8 +31,13 @@ def predict(input_df: pd.DataFrame, tfidf_path:str , model_path: str):
     with open(model_path, "rb") as model_file:
         model = pickle.load(model_file)
 
+    # Load vocabulary
+    with open(tfidf_voc_path, "rb") as f:
+        vocab = pickle.load(f)
+
     tfidf_vectorizer = TfidfVectorizer(**params)
     tfidf_vectorizer.set_params(preprocessor=process_text)
+    tfidf_vectorizer.set_params(vocabulary=vocab)
 
     # Transform the input text using the TF-IDF vectorizer
     text_data = input_df.to_pandas()["quote"]