Clean state

app.py

@@ -35,9 +35,12 @@ cmap = plt.get_cmap('RdYlGn_r')  # Red-Yellow-Green colormap, reversed
 def count_string_tokens(string: str, model: str) -> int:
     try:
         encoding = tiktoken.encoding_for_model(model.split('/')[-1])
-    except KeyError:
-        print(f"Model {model} not found. Using cl100k_base encoding.")
-        encoding = tiktoken.get_encoding("cl100k_base")
+    except:
+        try:
+            encoding = tiktoken.encoding_for_model(model.split('/')[-2] + '/' + model.split('/')[-1])
+        except KeyError:
+            print(f"Model {model} not found. Using cl100k_base encoding.")
+            encoding = tiktoken.get_encoding("cl100k_base")
     return len(encoding.encode(string))
 
 def calculate_total_cost(prompt_tokens: int, completion_tokens: int, model: str) -> float:
@@ -85,26 +88,42 @@ def compute_all(input_type, prompt_text, completion_text, prompt_tokens, complet
         df[col] = df[col].str.replace('$', '').astype(float)
 
     if len(df) > 1:
-        def apply_color(val, min, max):
-            norm = plt.Normalize(min, max)
+        norm = plt.Normalize(df['Total Cost'].min(), df['Total Cost'].max())
+        
+        def apply_color(val):
             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"])
+        
     else:
-        df = df.style.applymap(lambda x: 'background-color: rgba(0, 0, 0, 0)', 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',
-        'vertical-align': 'middle'
-    })
-    df = df.set_properties(**{'font-weight': 'bold'}, subset=['Total Cost'])
-    return df
+        def apply_color(val):
+            return "background-color: rgba(0, 0, 0, 0)"
+    
+    
+    # Apply colors and formatting
+    def style_cell(val, column):
+        style = ''
+        if column == 'Total Cost':
+            style += 'font-weight: bold; '
+        if column in ['Prompt Cost', 'Completion Cost', 'Total Cost']:
+            val = f'${val:.6f}'
+        if column == 'Total Cost':
+            style += apply_color(float(val.replace('$', '')))
+        return f'<td style="{style}">{val}</td>'
+
+    html_table = '<table class="styled-table">'
+    html_table += '<thead><tr>'
+    for col in df.columns:
+        html_table += f'<th>{col}</th>'
+    html_table += '</tr></thead><tbody>'
+    for _, row in df.iterrows():
+        html_table += '<tr>'
+        for col in df.columns:
+            html_table += style_cell(row[col], col)
+        html_table += '</tr>'
+    html_table += '</tbody></table>'
+    
+    return html_table
 
 def toggle_input_visibility(choice):
     return (
@@ -112,7 +131,25 @@ def toggle_input_visibility(choice):
         gr.Group(visible=(choice == "Token Count Input"))
     )
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, secondary_hue=gr.themes.colors.orange)) as demo:
+with gr.Blocks(css="""
+    .styled-table {
+        border-collapse: collapse;
+        margin: 25px 0;
+        font-family: Arial, sans-serif;
+        width: 100%;
+    }
+    .styled-table th, .styled-table td {
+        padding: 12px 15px;
+        text-align: left;
+        vertical-align: middle;
+    }
+    .styled-table tbody tr {
+        border-bottom: 1px solid #dddddd;
+    }
+    .styled-table tbody tr:nth-of-type(even) {
+        background-color: #f3f3f3;
+    }
+""", 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 prices data from [BerriAI's litellm](https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json).
@@ -142,12 +179,12 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.yellow, seconda
                     max_price = gr.Slider(label="Max Price per Input Token", minimum=0, maximum=0.001, step=0.00001, value=0.001)
                 litellm_provider = gr.Dropdown(label="Inference Provider", choices=["Any"] + TOKEN_COSTS['litellm_provider'].unique().tolist(), value="Any")
         
-            model = gr.Dropdown(label="Models (can select multiple)", choices=TOKEN_COSTS['model'].tolist(), value="anyscale/meta-llama/Meta-Llama-3-8B-Instruct", multiselect=True)
+            model = gr.Dropdown(label="Models (at least 1)", choices=TOKEN_COSTS['model'].tolist(), value="anyscale/meta-llama/Meta-Llama-3-8B-Instruct", multiselect=True)
         
-    gr.Markdown("## ➡️ Resulting Costs:")
+    gr.Markdown("## Resulting Costs 👇")
 
     with gr.Row():
-        results_table = gr.Dataframe()
+        results_table = gr.HTML()
 
     input_type.change(
         toggle_input_visibility,