utils/filter_utils.py

# Utility functions for filtering the dataframe

import pandas as pd

def filter_cols(df):

    df = df[[
    'Model Name', 
    'Input $/1M', 
    'Output $/1M',
    'Release Date', 
    'Context Size', 
    'Average Clemscore',
    'Average Latency (s)',
    'Parameter Size (B)',
    'License'
    ]]

    return df


def filter(df, language_list, parameters, input_price, output_price, multimodal,
           context, open_weight, start, end, license ):
    
    if not df.empty:  # Check if df is non-empty
        df = df[df['Languages'].apply(lambda x: all(lang in x for lang in language_list))]
    
    if not df.empty:  # Check if df is non-empty
        df = df[(df['Parameter Size (B)'] >= pow(2, parameters[0])) & (df['Parameter Size (B)'] <= pow(2, parameters[1]))]
    
    if not df.empty:  # Check if df is non-empty
        df = df[(df['Input $/1M'] >= input_price[0]) & (df['Input $/1M'] <= input_price[1])]
    
    if not df.empty:  # Check if df is non-empty
        df = df[(df['Output $/1M'] >= output_price[0]) & (df['Output $/1M'] <= output_price[1])]

    if not df.empty:  # Check if df is non-empty
        if "Image" in multimodal:
            df = df[df['Multimodality Image'] == True]
        if "Multi-Image" in multimodal:
            df = df[df['Multimodality Multiple Image'] == True]
        if "Audio" in multimodal:
            df = df[df['Multimodality Audio'] == True]
        if "Video" in multimodal:
            df = df[df['Multimodality Video'] == True]

    if not df.empty:  # Check if df is non-empty
        df = df[(df['Context Size'] >= (2**context[0])*1024) & (df['Context Size'] <= (2**context[1])*1024)]

    if not df.empty:  # Check if df is non-empty
        if "Open" in open_weight and "Commercial" not in open_weight:
            df = df[df['Open Weight'] == True]
        elif "Commercial" in open_weight and "Open" not in open_weight:
            df = df[df['Open Weight'] == False]
        
    if not df.empty:  # Check if df is non-empty
        df = df[df['License Name'].apply(lambda x: any(lic in x for lic in license))]

    # Convert 'Release Date' to int temporarily
    if not df.empty:  # Check if df is non-empty
        df['Temp Date'] = pd.to_datetime(df['Temp Date']).astype(int) // 10**9  # Convert to seconds since epoch

    # Convert start and end to int (seconds since epoch)
    start = int(pd.to_datetime(start).timestamp())  
    end = int(pd.to_datetime(end).timestamp())    

    # Filter based on the converted 'Release Date'
    if not df.empty:  # Check if df is non-empty
        df = df[(df['Temp Date'] >= start) & (df['Temp Date'] <= end)]

    df = filter_cols(df)
    return df  # Return the filtered dataframe