sentiment_labeling.py

import pandas as pd

def label_sentiment(rating):
    """Label sentiment based on the rating."""
    if rating in [1, 2]:
        return 'negative'
    elif rating == 3:
        return 'neutral'
    elif rating in [4, 5]:
        return 'positive'
    else:
        return 'unknown'

def add_sentiment_column(data: pd.DataFrame, rating_col_name='rating'):
    """Add a sentiment column to the dataframe based on the ratings."""
    if rating_col_name not in data.columns:
        raise ValueError(f"Column '{rating_col_name}' not found in the dataframe.")
    
    data['sentiment'] = data[rating_col_name].apply(label_sentiment)
    return data