src/prediction.py

import pandas as pd

def predict(data, model, prep):
    df = pd.DataFrame([data])

    # Handle country
    top_country = ['PRT', 'GBR', 'FRA', 'ESP', 'DEU', 'ITA', 'IRL', 'BEL', 'BRA', 'NLD']
    df['country'] = df['country'].apply(lambda x: x if x in top_country else 'Other')

    # Fix missing columns
    if 'agent' not in df.columns:
        df['agent'] = 0

    if 'company' not in df.columns:
        df['company'] = 0

    # Transform
    df_processed = prep.transform(df)

    # Predict
    result = model.predict(df_processed)

    return result[0]