plotting perplexity vs text length with text hover

app.py

@@ -1,32 +1,48 @@
 import dash
-from dash.exceptions import PreventUpdate
+import plotly.express as px
 from dash import dcc, html
 from dash.dependencies import Input, Output
-import plotly.express as px
-import pandas as pd
+from dash.exceptions import PreventUpdate
+from datasets import load_dataset
 
 # Create dash app
 app = dash.Dash(__name__)
 
-# Set dog and cat images
-dogImage = "https://www.iconexperience.com/_img/v_collection_png/256x256/shadow/dog.png"
-catImage = "https://d2ph5fj80uercy.cloudfront.net/06/cat3602.jpg"
-
-# Generate dataframe
-df = pd.DataFrame(
-    dict(
-        x=[1, 2],
-        y=[2, 4],
-        images=[dogImage, catImage],
-    )
-)
 
-# Create scatter plot with x and y coordinates
-fig = px.scatter(df, x="x", y="y", custom_data=["images"])
+def get_dataset(name, n_items=1000):
+    ola_path = f"ola13/small-{name}-dedup"
+    dataset = load_dataset(ola_path, split="train").shuffle().select(range(n_items)).to_pandas()
+    dataset["text_length"] = dataset.apply(lambda doc: len(doc["text"]), axis=1)
+
+    for column in dataset.columns:
+        if column not in ["text", "perplexity", "text_length"]:
+            dataset = dataset.drop(column, axis=1)
 
+    dataset = dataset.sort_values("perplexity")
+
+    max_perp = dataset["perplexity"].max()
+    return dataset, max_perp
+
+
+# names = ["oscar", "the_pile", "c4", "roots_en"]
+name = "oscar"
+df, max_perplexity = get_dataset(name)
+
+# Create scatter plot with x and y coordinates
+fig = px.scatter(df, x="perplexity", y="text_length", custom_data=["text"])
 # Update layout and update traces
 fig.update_layout(clickmode='event+select')
-fig.update_traces(marker_size=20)
+fig.update_traces(marker_size=3)
+fig.update_xaxes(title_text="Perplexity (log scale)", type="log")
+fig.update_yaxes(title_text="Text Length (log scale)", type="log")
+
+styles = {
+    'textbox': {
+        'border': 'thin lightgrey solid',
+        'overflowX': 'scroll',
+        "whiteSpace": "pre-wrap;"
+    }
+}
 
 # Create app layout to show dash graph
 app.layout = html.Div(
@@ -35,14 +51,14 @@ app.layout = html.Div(
             id="graph_interaction",
             figure=fig,
         ),
-        html.Img(id='image', src='')
+        html.Div(id='text', style=styles['textbox'])
     ]
 )
 
 
 # html callback function to hover the data on specific coordinates
 @app.callback(
-    Output('image', 'src'),
+    Output('text', 'children'),
     Input('graph_interaction', 'hoverData'))
 def open_url(hoverData):
     if hoverData:
@@ -52,4 +68,4 @@ def open_url(hoverData):
 
 
 if __name__ == '__main__':
-    app.run_server(port=7860, host="0.0.0.0", debug=True, )
+    app.run_server(port=7860, host="0.0.0.0", debug=True)