add cluster

data/clustering/info.csv

@@ -0,0 +1,106 @@
+,cluster_id,cluster_summaries,cluster_position_x,cluster_position_y
+0,-1,None,9.926462,4.7121987
+1,0,Philosophical/Spiritual Introspection,10.312462,1.2666532
+2,1,"Scholarships",8.167274,4.8995786
+3,2,Politics,8.81142,2.4859838
+4,3,Theology,9.615214,0.3783942
+5,4,Dating,4.985182,1.8439052
+6,5,Accommodation,11.457769,5.080919
+7,6,Football,6.6154537,-1.6859366
+8,7,Film Festival,6.9734483,1.4548192
+9,8,Culinary,13.426296,4.5412893
+10,9,Music,6.0653744,0.7536916
+11,10,Gambling,3.124241,3.2533677
+12,11,Baseball,7.133596,-2.4256644
+13,12,Technology,6.4929094,6.768577
+14,13,Website Policies,4.873843,5.771508
+15,14,Weddings,11.815845,3.7894728
+16,15,Gaming,5.529167,2.9530518
+17,16,Commodities/Services Provision,10.453564,5.8489122
+18,17,Crafts,13.287651,6.4237967
+19,18,Automobiles,9.9531145,8.840178
+20,19,Watches,13.893139,9.859185
+21,20,Dogs,12.595798,3.5351615
+22,21,Photography,10.7942295,3.5504062
+23,22,Legalities,8.942016,4.72733
+24,23,Consumer Electronics,7.078649,8.338984
+25,24,Insulation,10.520957,7.914946
+26,25,Cannabis,14.317424,3.2114828
+27,26,Footwear,15.052116,7.6956415
+28,27,Real Estate,9.536316,6.103533
+29,28,Relocation,10.205071,7.1883316
+30,29,Sports betting,3.2779586,2.443366
+31,30,Narratives,7.613535,1.8300554
+32,31,Dating,4.788838,2.1900373
+33,32,Apparel/Clothing,14.394226,7.3073387
+34,33,User Authentication,5.265638,6.4014487
+35,34,Academicwriting,6.9187264,3.4357684
+36,35,Sports,7.4969172,-2.086585
+37,36,Fashion/Lifestyle Products,13.821669,7.7150764
+38,37,Diverse events,9.437052,2.2438836
+39,38,Blockchain/Cryptocurrency,7.7586045,6.9439344
+40,39,Online Businesses/Marketing,6.522259,5.219268
+41,40,Healthcare,11.425277,2.3801014
+42,41,Home Decor,12.878046,7.2632184
+43,42,Biomedicine,12.789575,2.3376262
+44,43,Jewelry,14.259997,8.653363
+45,44,Addiction,11.561383,1.3774762
+46,45,Products,11.711758,8.423251
+47,46,Multi-purposefulness,11.080702,7.4574013
+48,47,"Mass transit",9.910158,5.4402313
+49,48,Ethernet,6.9763823,7.7909245
+50,49,Legal,9.516912,4.636553
+51,50,E-commerce,13.263438,8.6548195
+52,51,Audio,7.717162,8.903019
+53,52,Infrastructure,10.52904,5.369669
+54,53,Firearms,11.062812,9.268473
+55,54,Freight/Logistics,9.551044,7.0336204
+56,55,Products,12.073747,7.645973
+57,56,Vaccinations,11.9387045,2.7824683
+58,57,Artwork,11.019163,4.1677165
+59,58,Viticulture,14.223523,5.0761614
+60,59,WordPress,5.9597983,5.824579
+61,60,Cosmetics/Dermatology,15.093273,3.4669027
+62,61,Software,6.375921,6.4298844
+63,62,Dentistry,14.76626,1.1620314
+64,63,Pest Control,13.201735,3.6806118
+65,64,SEO,5.720493,5.238112
+66,65,Lottery,1.7142816,2.9782674
+67,66,Narratives,8.460977,1.0804662
+68,67,Waste Reduction & Recycling,10.634534,6.959523
+69,68,Communication,6.438943,5.9467845
+70,69,Orthopedics,13.005415,1.1908791
+71,70,Home Decor & Furniture,12.732457,7.876862
+72,71,Education,7.6568975,3.4944353
+73,72,Sports,7.295141,-0.7343214
+74,73,Social Media Advertising,6.133886,4.8547883
+75,74,Privacy,4.756733,6.3598356
+76,75,Website design,6.1168823,5.465095
+77,76,Roofing,11.389448,8.080609
+78,77,Nutrition/Supplements,13.631578,2.5334294
+79,78,Haircare/Hairstyling,15.544645,4.54254
+80,79,Cookies,4.341592,6.819268
+81,80,International Trade,8.993828,6.4757586
+82,81,Entrepreneurial Resources,9.435777,5.3340797
+83,82,Cricket,6.5171986,-1.245905
+84,83,Crafts,13.852216,7.049825
+85,84,Floristry,13.407425,5.8741536
+86,85,Genealogy,9.530803,1.6548243
+87,86,Mental Health,11.074349,1.6069281
+88,87,Volunteerism,10.145443,3.6734574
+89,88,Lighting,11.385381,8.93693
+90,89,Artificial Intelligence,6.5306387,6.2178063
+91,90,Business,7.471462,6.4142885
+92,91,E-commerce,13.638669,6.5098934
+93,92,Urbanization/Over-tourism,10.221115,6.100654
+94,93,Events,10.8449,3.9822264
+95,94,Pharmaceuticals/Biotechnology,12.318266,2.4331784
+96,95,Professional Wrestling,6.856304,-0.65598303
+97,96,Various,9.3211975,3.4894605
+98,97,Medicine,13.17882,2.1281319
+99,98,Community Engagement,9.848856,3.5187004
+100,99,Fitness,12.504849,0.9134393
+101,100,Bathroom Design & Toilet Engineering,11.779076,7.2920136
+102,101,Business Development,7.328447,5.659843
+103,102,Sports,7.6370654,-1.0701839
+104,103,Sexuality,13.817207,1.6510898

index.html

@@ -5,7 +5,9 @@
     <script src="https://cdnjs.cloudflare.com/ajax/libs/mathjs/12.4.2/math.min.js" charset="utf-8"></script>
     <script src="https://cdn.plot.ly/plotly-2.32.0.min.js" charset="utf-8"></script>
     <script src="https://cdnjs.cloudflare.com/ajax/libs/lodash.js/4.17.21/lodash.min.js" charset="utf-8"></script>
+    <script src="https://cdnjs.cloudflare.com/ajax/libs/PapaParse/5.4.1/papaparse.js" integrity="sha512-M0cjXJTonbWEdLI3HJIoJSQBb9980RWmOCk+tvWkhgFrAZqSSIg1+1Db/vDu7Qk9W3L90gBynve17PYvarjfQA==" crossorigin="anonymous" referrerpolicy="no-referrer"></script>
     <script type="module" src="src/plotting.js"></script>
+    <script type="module" src="src/clusters.js"></script>
     <link rel="stylesheet" href="style.css">
     <meta name="viewport" content="width=device-width, initial-scale=1">
     <meta charset="utf8">
@@ -159,9 +161,12 @@
     </script>
 </d-front-matter>
 <d-title>
-    <figure class="l-page">
-        <img src="banner.png" alt="FineWeb">
-    </figure>
+    <div class="main-plot-container l-page">
+        <figure>
+            <img src="banner.png" alt="FineWeb">
+        </figure>
+        <div id="clusters-plot"></div>
+    </div>
 </d-title>
 <d-byline></d-byline>
 <d-article>

src/clusters.js

@@ -0,0 +1,223 @@
+// plotly Dark24
+const COLORS = [
+
+    ['46', '145', '229'], 
+    ['225', '95', '153'], 
+    ['28', '167', '28'], 
+    ['251', '13', '13'], 
+    ['218', '22', '255'], 
+    // ['34', '42', '42'], Black, which makes the text unreadable
+    ['182', '129', '0'], 
+    ['117', '13', '134'], 
+    ['235', '102', '59'], 
+    ['81', '28', '251'], 
+    ['0', '160', '139'], 
+    ['251', '0', '209'], 
+    ['252', '0', '128'], 
+    ['178', '130', '141'], 
+    ['108', '124', '50'], 
+    ['119', '138', '174'], 
+    ['134', '42', '22'], 
+    ['167', '119', '241'], 
+    ['98', '0', '66'], 
+    ['22', '22', '167'], 
+    ['218', '96', '202'], 
+    ['108', '69', '22'], 
+    // ['13', '42', '99'], Black
+    ['175', '0', '56']
+]
+
+const getLabelHoverFormat = (row) => {
+    return `<b>Text</b>: ${row.text}<br><b>Edu label</b>: ${row.eduScore}`
+}
+
+// Number of annotations to display
+const K = 15
+
+
+function createLabelOrderMapping(labels) {
+    const labelCounts = labels.reduce((acc, label) => {
+        acc[label] = (acc[label] || 0) + 1;
+        return acc;
+    }, {});
+
+    const sortedLabels = Object.entries(labelCounts).sort((a, b) => b[1] - a[1]).map(entry => entry[0]);
+
+    const labelOrder = {};
+    sortedLabels.forEach((label, index) => {
+        labelOrder[label] = index;
+    });
+    return labelOrder;
+}
+
+
+
+
+const parseAnnotations = async (file) => {
+    return (await readCSV(file)).filter((cluster_summary) => {
+        return parseInt(cluster_summary.cluster_id) != -1
+    }).map((cluster_summary) => {
+        return {
+            x: parseFloat(cluster_summary.cluster_position_x),
+            y: parseFloat(cluster_summary.cluster_position_y),
+            label: parseInt(cluster_summary.cluster_id),
+            text: cluster_summary.cluster_summaries,
+        }
+    })
+}
+
+const addStylingToAnnotations = (annotations) => {
+    return annotations.map((annotation, i) => {
+        return {
+            ...annotation,
+            showarrow: false,
+            font: {
+                size: 14,
+                color: 'black',
+                weight: 'bold'
+              },
+            bgcolor: getColor(annotation.label, 0.9),
+            borderpad: 2, // Add padding around the text
+        }
+    })
+}
+
+const getRelevantAnnotations = (annotations, x0, x1, y0, y1, k=K) => {
+    const relevant_annotations = annotations.filter((annotation) => {
+        return annotation.x >= x0 && annotation.x <= x1 && annotation.y >= y0 && annotation.y <= y1
+    })
+    return relevant_annotations.sort((a, b) => a.ord - b.ord).slice(0, k);
+}
+
+
+const getMinMaxTracesArea = (traces) => {
+    const x0 = Math.min(...traces.map(trace => trace.x));
+    const x1 = Math.max(...traces.map(trace => trace.x));
+    const y0 = Math.min(...traces.map(trace => trace.y));
+    const y1 = Math.max(...traces.map(trace => trace.y));
+    return {x0, x1, y0, y1};
+}
+
+const readData = async () => {
+    return (await readCSV('data/clustering/data.csv')).map(row => ({
+        x: parseFloat(row.X),
+        y: parseFloat(row.Y),
+        eduScore: parseFloat(row.edu_labels),
+        label: parseInt(row.cluster_labels),
+        text: row.content_display,
+    }));
+}
+
+async function plotClusters() {
+    const parent = document.getElementById('clusters-plot');
+    const data = await readData();
+    const traces = [{
+            type: 'scatter',
+            mode: 'markers',
+            x: data.map(row => row.x),
+            y: data.map(row => row.y),
+            marker: {
+                color: data.map(row => getColor(row.label, 1.0)),
+                size: 5,
+                opacity: 8
+            },
+            hoverinfo: 'text',
+            hovertext: data.map(row => getLabelHoverFormat(row)),
+            hoverlabel: {
+                bgcolor: 'white'
+            },
+    }];
+    const labelOrder = createLabelOrderMapping(data.map(row => row.label));
+    const annotations = (await parseAnnotations('data/clustering/info.csv')).map(
+        (annot) => {
+            return {
+                ...annot,
+                ord: labelOrder[annot.label]
+            }
+        }
+    );
+
+    const {x0, x1, y0, y1} = getMinMaxTracesArea(data);
+    const layout = {
+        height: 550,
+        width: parent.clientWidth,
+        xaxis: {
+            showticklabels: false,
+            showgrid: false,
+            zeroline: false,
+            title: {
+                text: "Fineweb dataset (clustered using TODO and labeled using TODO),<br> zoom in to see more",
+                font: {
+                    size: 16,
+                    style: 'italic'
+                },
+                // italic
+            },
+        },
+        yaxis: {
+            showticklabels: false,
+            showgrid: false,
+            zeroline: false,
+        },
+        annotations: addStylingToAnnotations(getRelevantAnnotations(annotations, x0, x1, y0, y1)),
+        font: {
+            family: "apple-system, Arial, sans-serif",
+        },
+        margin: {
+            t: 0,
+            b: 30,
+        },
+    };
+
+    Plotly.newPlot(parent, traces, layout);
+
+    parent.on('plotly_relayout', (eventdata) => {
+        // First option zoomed in
+        if (eventdata["xaxis.range[0]"]) {
+            const [x0, x1] = [eventdata['xaxis.range[0]'], eventdata['xaxis.range[1]']];
+            const [y0, y1] = [eventdata['yaxis.range[0]'], eventdata['yaxis.range[1]']];
+            // Idk maybe we can even recompute the ordering, but I think it's fine to use the global one
+            const relevant_annotations = getRelevantAnnotations(annotations, x0, x1, y0, y1);
+            Plotly.relayout(parent, {...layout, annotations: addStylingToAnnotations(relevant_annotations)});
+        }
+        // Zoom reset
+        else if (eventdata["xaxis.autorange"]){
+            const {x0, x1, y0, y1} = getMinMaxTracesArea(data);
+            const relevant_annotations = getRelevantAnnotations(annotations, x0, x1, y0, y1);
+            Plotly.relayout(parent, {...layout, annotations: addStylingToAnnotations(relevant_annotations)});
+        }
+        // Otherwise it's just the relayout itself
+    });
+   
+    window.addEventListener("resize", () => {
+        // If the window size is smaller than 768, we don't care as it's not shown
+        if (window.innerWidth < 768) {
+            return;
+        }
+        Plotly.relayout(parent, {
+            width: parent.offsetWidth,
+        });
+    });
+}
+
+document.addEventListener("DOMContentLoaded", () => {
+    plotClusters();
+});
+
+
+const readCSV = async (file) => {
+    const data = await fetch(file)
+    const text = await data.text()
+    const csv = Papa.parse(text, {header: true});
+    return csv.data;
+}
+
+
+
+const getColor = (i, opacity) => {
+    if (i < 0) {
+        i = i * -1
+    }
+    console.log(COLORS[i % COLORS.length])
+    return `rgba(${COLORS[i % COLORS.length].join(',')}, ${opacity})`
+}

src/plotting.js

@@ -1,6 +1,6 @@
 const TASK_ID_TO_NAME = {
     // Ablations
-    "agg_score": "Aggregate Score",
+    agg_score: "Aggregate Score",
     "commonsense_qa/acc_norm": "Commonsense QA Norm",
     "hellaswag/acc_norm": "HellaSwag",
     "openbookqa/acc_norm": "OpenBook QA Norm",
@@ -11,27 +11,26 @@ const TASK_ID_TO_NAME = {
     "mmlu/acc_norm": "MMLU",
 
     // Stats
-    "ccnet": "CCNet",
+    ccnet: "CCNet",
 };
 
 const DATASET_ID_TO_NAME = {
-    "pii_removed": "Fineweb",
-    "allenai_c4_en": "C4",
+    pii_removed: "Fineweb",
+    allenai_c4_en: "C4",
     "tiiuae_falcon-refinedweb_data": "RefinedWeb",
     "red-pajama-v2_jsonl-deduplicated-extract": "RedPajamaV2",
     "dolma-sample": "Dolma1.6",
-    "dedup_minhash_independent_output": "Individual Dedup MinHash",
+    dedup_minhash_independent_output: "Individual Dedup MinHash",
 };
 
 const DEFAULT_SETTINGS = {
-    "slider": {
-        "max": 30,
-        "min": 0,
-        "default": 0
+    slider: {
+        max: 30,
+        min: 0,
+        default: 0,
     },
-    "defaultMetric": "agg_score",
-}
-
+    defaultMetric: "agg_score",
+};
 
 const DEFAULT_LAYOUT = {
     title: {
@@ -99,44 +98,43 @@ const getAutoRange = (traces) => {
     let minX = Math.min(...traces.flatMap((trace) => trace.x));
     let maxX = Math.max(...traces.flatMap((trace) => trace.x));
     return [minX * 0.95, maxX * 1.05];
-}
+};
 
 const init_ablation_plot = function () {
     const plotElements = document.querySelectorAll('[id^="plot-"]');
-
-
-
-        plotElements.forEach(async (plotElement) => {
-            const plotName = plotElement.id.replace('plot-', '');
-            const indexData = await fetch(`data/plots/${plotName}/index.json`).then(
-                (response) => response.json()
-            );
-            const settings = _.merge(DEFAULT_SETTINGS, indexData.settings);
-            const indexMapping = indexData.files;
-            const { dropdown, slider, plot } = createAblationPlottingElements(
-                plotElement,
-                indexMapping,
-                settings
-            );
-            plot.id = `graph-${plotName}`;
-            dropdown.addEventListener("change", () => updatePlot(dropdown, slider));
-            let timeoutId;
-            // Debounce the slider
-            if (slider !== undefined) {
-                slider.addEventListener("input", () => {
-                    clearTimeout(timeoutId);
+    plotElements.forEach(async (plotElement) => {
+        const plotName = plotElement.id.replace("plot-", "");
+        const indexData = await fetch(`data/plots/${plotName}/index.json`).then(
+            (response) => response.json()
+        );
+        const settings = _.merge({}, DEFAULT_SETTINGS, indexData.settings);
+        const indexMapping = indexData.files;
+        const { dropdown, slider, plot } = createAblationPlottingElements(
+            plotElement,
+            indexMapping,
+            settings
+        );
+        plot.id = `graph-${plotName}`;
+        dropdown.addEventListener("change", () => updatePlot(dropdown, slider));
+        let timeoutId;
+        // Debounce the slider
+        if (slider !== undefined) {
+            slider.addEventListener("input", () => {
+                clearTimeout(timeoutId);
                 timeoutId = setTimeout(() => {
-                        updatePlot(dropdown, slider);
-                    }, 500);
-                });
-            }
+                    updatePlot(dropdown, slider);
+                }, 500);
+            });
+        }
         // Shared plot
-        Plotly.newPlot(plot, DEFAULT_LAYOUT);
+        Plotly.newPlot(plot, []);
 
         async function updatePlot(dropdown, slider) {
             const metricName = dropdown.value;
             const sliderValue = parseInt(slider?.value ?? 0);
-            const metricData = await fetch(`data/plots/${plotName}/${indexMapping[metricName]["file"]}`).then((response) => response.json());
+            const metricData = await fetch(
+                `data/plots/${plotName}/${indexMapping[metricName]["file"]}`
+            ).then((response) => response.json());
             const traces = metricData?.traces?.[metricName] ?? [];
             for (const key in metricData?.data ?? []) {
                 const traceData = metricData.data[key];
@@ -161,7 +159,11 @@ const init_ablation_plot = function () {
                 {
                     width: width,
                     yaxis: { title: { text: TASK_ID_TO_NAME[metricName] } },
-                    xaxis: { range: settings.autoSetXRange ? getAutoRange(traces) : undefined },
+                    xaxis: {
+                        range: settings.autoSetXRange
+                            ? getAutoRange(traces)
+                            : undefined,
+                    },
                 },
                 metricData.layout
             );
@@ -197,7 +199,9 @@ const createAblationPlottingElements = (
     plotElement.appendChild(plot);
     plotElement.appendChild(controls);
 
-    const metricOptions = Object.keys(indexMapping).filter(metric => metric in TASK_ID_TO_NAME);
+    const metricOptions = Object.keys(indexMapping).filter(
+        (metric) => metric in TASK_ID_TO_NAME
+    );
     // Dropdown
     const dropdownLabel = document.createElement("label");
     dropdownLabel.textContent = "Metric:";
@@ -216,7 +220,7 @@ const createAblationPlottingElements = (
     dropdownContainer.appendChild(dropdown);
     controls.appendChild(dropdownContainer);
 
-    let slider = undefined
+    let slider = undefined;
     if (settings.slider !== null) {
         const sliderLabel = document.createElement("label");
         sliderLabel.textContent = "Rolling window:";