tasks/utils/predict.py

import pickle
import re
import string
import pandas as pd
import sys
sys.path.append(".") 
from tasks.utils.preprocessing import process_text
    
def predict(input_df: pd.DataFrame, tfidf_path: str, model_path: str):
    """
    Predict the output using a saved TF-IDF vectorizer and Random Forest model.

    Parameters:
        input_df (pd.DataFrame): Input dataframe containing the text data.
        tfidf_path (str): Path to the saved TF-IDF vectorizer pickle file.
        model_path (str): Path to the saved Random Forest model pickle file.
        text_column (str): The name of the column in the dataframe containing the text data.

    Returns:
        pd.Series: Predictions for each row in the input dataframe.
    """
    # Load the TF-IDF vectorizer
    with open(tfidf_path, "rb") as tfidf_file:
        tfidf_vectorizer = pickle.load(tfidf_file)

    # Load the Random Forest model
    with open(model_path, "rb") as model_file:
        model = pickle.load(model_file)

    # Transform the input text using the TF-IDF vectorizer
    text_data = input_df.to_pandas()["quote"]
    text_features = tfidf_vectorizer.transform(text_data)

    # Make predictions using the loaded model
    predictions = model.predict(text_features)

    return predictions