Add table layout with multiple models

app.py

@@ -3,6 +3,7 @@ import pandas as pd
 import requests
 import json
 import tiktoken
+import matplotlib.pyplot as plt
 
 PRICES_URL = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
 
@@ -20,9 +21,16 @@ except Exception as e:
 
 TOKEN_COSTS = pd.DataFrame.from_dict(TOKEN_COSTS, orient='index').reset_index()
 TOKEN_COSTS.columns = ['model'] + list(TOKEN_COSTS.columns[1:])
-TOKEN_COSTS = TOKEN_COSTS.loc[~TOKEN_COSTS["model"].str.contains("sample_spec")]
-TOKEN_COSTS = TOKEN_COSTS.loc[~TOKEN_COSTS["input_cost_per_token"].isnull()]
+TOKEN_COSTS = TOKEN_COSTS.loc[
+    (~TOKEN_COSTS["model"].str.contains("sample_spec"))
+    & (~TOKEN_COSTS["input_cost_per_token"].isnull())
+    & (~TOKEN_COSTS["output_cost_per_token"].isnull())
+    & (TOKEN_COSTS["input_cost_per_token"] > 0)
+    & (TOKEN_COSTS["output_cost_per_token"] > 0)
+]
+TOKEN_COSTS["supports_vision"] = TOKEN_COSTS["supports_vision"].fillna(False)
 
+cmap = plt.get_cmap('RdYlGn_r')  # Red-Yellow-Green colormap, reversed
 
 def count_string_tokens(string: str, model: str) -> int:
     try:
@@ -50,27 +58,54 @@ def update_model_list(function_calling, litellm_provider, max_price, supports_vi
     list_models = filtered_models['model'].tolist()
     return gr.Dropdown(choices=list_models, value=list_models[0] if list_models else "No model found for this combination!")
 
-def compute_all(input_type, prompt_text, completion_text, prompt_tokens, completion_tokens, model):
-    if input_type == "Text Input":
-        prompt_tokens = count_string_tokens(prompt_text, model)
-        completion_tokens = count_string_tokens(completion_text, model)
-    else:  # Token Count Input
-        prompt_tokens = int(prompt_tokens * 1000)
-        completion_tokens = int(completion_tokens * 1000)
-
-    prompt_cost, completion_cost = calculate_total_cost(prompt_tokens, completion_tokens, model)
-    total_cost = prompt_cost + completion_cost
+def compute_all(input_type, prompt_text, completion_text, prompt_tokens, completion_tokens, models):
+    results = []
+    for model in models:
+        if input_type == "Text Input":
+            prompt_tokens = count_string_tokens(prompt_text, model)
+            completion_tokens = count_string_tokens(completion_text, model)
+        else:  # Token Count Input
+            prompt_tokens = int(prompt_tokens * 1000)
+            completion_tokens = int(completion_tokens * 1000)
+
+        prompt_cost, completion_cost = calculate_total_cost(prompt_tokens, completion_tokens, model)
+        total_cost = prompt_cost + completion_cost
+        
+        results.append({
+            "Model": model,
+            "Prompt Cost": f"${prompt_cost:.6f}",
+            "Completion Cost": f"${completion_cost:.6f}",
+            "Total Cost": f"${total_cost:.6f}"
+        })
     
-    return (
-        f"${prompt_cost:.6f}",
-        f"${completion_cost:.6f}",
-        f"${total_cost:.6f}"
-    )
+    df = pd.DataFrame(results)
+    
+    # Convert cost columns to numeric, removing the '$' sign
+    for col in ["Prompt Cost", "Completion Cost", "Total Cost"]:
+        df[col] = df[col].str.replace('$', '').astype(float)
+
+    if len(df) > 1:
+        def apply_color(val, min, max):
+            norm = plt.Normalize(min, max)
+            color = cmap(norm(val))
+            rgba = tuple(int(x * 255) for x in color[:3]) + (0.5,)
+            rgba = tuple(int(x * 255) for x in color[:3]) + (0.5,)  # 0.5 for 50% opacity
+            return f'background-color: rgba{rgba}'
+
+        min, max = df["Total Cost"].min(), df["Total Cost"].max()
+        df = df.style.applymap(lambda x: apply_color(x, min, max), subset=["Total Cost"])
+        df = df.format({"Prompt Cost": "${:.6f}", "Completion Cost": "${:.6f}", "Total Cost": "${:.6f}"})
+        df = df.set_properties(**{
+            'font-family': 'Arial, sans-serif',
+            'white-space': 'pre-wrap'
+        })
+        df = df.set_properties(**{'font-weight': 'bold'}, subset=['Total Cost'])
+    return df
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, secondary_hue=gr.themes.colors.orange)) as demo:
     gr.Markdown("""
     # Text-to-$$$: Calculate the price of your LLM runs
-    Based on data from [litellm](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json).
+    Based on prices data from [BerriAI's litellm](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json).
     """)
     
     with gr.Row():
@@ -91,19 +126,16 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
                 with gr.Column():
                     function_calling = gr.Checkbox(label="Supports Tool Calling", value=False)
                     supports_vision = gr.Checkbox(label="Supports Vision", value=False)
-                litellm_provider = gr.Dropdown(label="LiteLLM Provider", choices=["Any"] + TOKEN_COSTS['litellm_provider'].unique().tolist(), value="Any")
+                litellm_provider = gr.Dropdown(label="Inference Provider", choices=["Any"] + TOKEN_COSTS['litellm_provider'].unique().tolist(), value="Any")
             
             max_price = gr.Slider(label="Max Price per Token (input + output)", minimum=0, maximum=0.001, step=0.00001, value=0.001)
             
-            model = gr.Dropdown(label="Model", choices=TOKEN_COSTS['model'].tolist(), value=TOKEN_COSTS['model'].tolist()[0])
+            model = gr.Dropdown(label="Models (can select multiple)", choices=TOKEN_COSTS['model'].tolist(), value=[TOKEN_COSTS['model'].tolist()[0]], multiselect=True)
             
-            compute_button = gr.Button("Compute Costs", variant="secondary")
+            compute_button = gr.Button("Compute Costs ⚙️", variant="secondary")
 
-        with gr.Column(scale=1):
-            with gr.Group():
-                prompt_cost = gr.Textbox(label="Prompt Cost", interactive=False)
-                completion_cost = gr.Textbox(label="Completion Cost", interactive=False)
-                total_cost = gr.Textbox(label="Total Cost", interactive=False)
+        with gr.Column(scale=2):
+            results_table = gr.Dataframe(label="Cost Results")
 
     def toggle_input_visibility(choice):
         return (
@@ -134,7 +166,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
             completion_tokens_input,
             model
         ],
-        outputs=[prompt_cost, completion_cost, total_cost]
+        outputs=[results_table]
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -1,2 +1,3 @@
 pandas
-tiktoken
+tiktoken
+matplotlib