from fasthtml.common import *
from fasthtml.components import *
from fasthtml.components import (
    D_title,
    D_article,
    D_front_matter,
    D_contents,
    D_byline,
    D_bibliography,
    D_appendix,
    D_cite,
)
from plotly import graph_objects as go
from fh_plotly import plotly2fasthtml
import pandas as pd
import json
from rich import print
import overview
import curated
import web
import common
import results
from pybtex.database import parse_file
import data_viewer


app, rt = fast_app(
    debug=True,
    pico=False,
    hdrs=(
        Meta(charset="UTF-8"),
        Meta(name="viewport", content="width=device-width, initial-scale=1.0"),
        Script(src="https://distill.pub/template.v2.js"),
        Script(src="https://unpkg.com/htmx.org@next/dist/htmx.min.js"),
        Script(src="https://cdn.plot.ly/plotly-latest.min.js"),
        Link(rel="stylesheet", href="style.css"),
        MarkdownJS(),
    ),
)


front_matter = """
<d-front-matter>
<script id='distill-front-matter' type="text/json">{
    "title": "",
    "description": "",
    "published": "",
    "affiliation": {},
    "authors": [
      {
        "author":"",
        "authorURL":""
      }
    ],
    "katex": {
      "delimiters": [
        {"left": "$$", "right": "$$", "display": false}
      ]
    }
  }
</script>
</d-front-matter>
"""


def read_bibs():
    bib_data = parse_file("bibliography.bib")
    cits = []
    for key in bib_data.entries.keys():
        cits.append(D_cite(bibtex_key=key))
    return cits


@app.get("/bibliography.bib")
def get():
    return FileResponse("bibliography.bib")


@app.get("/")
def main():
    return Div(
        D_title(
            H1(
                "TxT360: a globally deduplicated dataset for LLM pretraining",
                cls="l-body",
                style="text-align: center;",
            ),
            Div(
                Img(src="images/llm360_logo.png"),
                id="title-plot",
                cls="main-plot-container l-page",
            ),
        ),
        Div(D_byline(), NotStr(front_matter), style="display: none;"),
        D_article(
            D_contents(
                Nav(
                    H3("Table of Contents"),
                    Div(
                        A(
                            "TxT360",
                            href="/intro#section1",
                            hx_get="/intro#section1",
                            hx_target="#inner-text",
                        )
                    ),
                    Div(
                        Ul(
                            Li(
                                A(
                                    "About TxT360",
                                    href="/intro#section1",
                                    hx_get="/intro#section1",
                                    hx_target="#inner-text",
                                )
                            ),
                            Li(
                                A(
                                    "Motivation Behind Txt360",
                                    href="/intro#section2",
                                    hx_get="/intro#section2",
                                    hx_target="#inner-text",
                                )
                            ),
                            Li(
                                A(
                                    "Generalizable Approach to Data Processing",
                                    href="/intro#section3",
                                    hx_get="/intro#section3",
                                    hx_target="#inner-text",
                                )
                            ),
                        ),
                    ),
                    Div(
                        A(
                            "Global Processing Steps",
                            href="/common#section1",
                            hx_get="/common#section1",
                            hx_target="#inner-text",
                        )
                    ),
                    Div(
                        A(
                            "Web Data Processing",
                            href="/webdata",
                            hx_get="/webdata",
                            hx_target="#inner-text",
                        )
                    ),
                    Div(
                        A(
                            "Curated Sources Processing",
                            href="/curated",
                            hx_get="/curated",
                            hx_target="#inner-text",
                        )
                    ),
                    Div(
                        A(
                            "TxT360 Results",
                            href="/results",
                            hx_get="/results",
                            hx_target="#inner-text",
                        ),
                    ),
                    role="navigation",
                    cls="l-text figcaption",
                ),
            ),
            intro(),
        ),
        D_appendix(D_bibliography(src="bibliography.bib")),
        Div(*read_bibs(), style="display: none;"),
    )


dataset_comparison1 = pd.DataFrame(
    {
        "Dataset": [
            "TxT360",
            "FineWeb",
            "RefinedWeb",
            "RedPajama-v2",
            "C4",
            "Dolma",
            "RedPajama-v1",
            "The Pile",
        ],
        "CommonCrawl": [
            "99 Snapshots",
            "96 Snapshots",
            "90 Snapshots",
            "84 Snapshots",
            "1 Snapshots",
            "24 Snapshots",
            "5 Snapshots",
            "0.6% of 74 Snapshots",
        ],
        "Papers": [
            "5 Sources",
            "-",
            "-",
            "-",
            "-",
            "1 Source",
            "1 Source",
            "4 Sources",
        ],
        "Wikipedia": [
            "310+ Languages",
            "-",
            "-",
            "-",
            "-",
            "what does a check mark mean?",
            "what does a check mark mean?",
            "English Only",
        ],
        "FreeLaw": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
        "DM Math": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
        "USPTO": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
    }
)

# Apply table styling: Light green for the header, alternating white and light grey for rows
styled_table = (
    dataset_comparison1.style.set_properties(
        **{"background-color": "#E1EEDB"},
        subset=pd.IndexSlice[0, :],  # Row 0 with a light green background
    )
    .apply(
        lambda x: [
            "background-color: #E1EEDB"
            if i == 0
            else (
                "background-color: rgb(237, 242, 251)"
                if i % 2 == 0
                else "background-color: white"
            )
            for i in range(len(x))
        ],
        axis=0,
    )
    .hide(axis="index")
)  # Hide the row index

# Use _repr_html_() method to get the HTML representation of the styled DataFrame
table_html = styled_table._repr_html_()
# table_html = dataset_comparison1.to_html(index=False, border=0)
table_div_1 = Div(NotStr(table_html), style="margin: 40px;")

dataset_comparison2 = pd.DataFrame(
    {
        "Dataset": [
            "TxT360",
            "FineWeb",
            "RefinedWeb",
            "RedPajama-v2",
            "C4",
            "Dolma",
            "RedPajama-v1",
            "The Pile",
        ],
        "PG-19": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "Included",
            "Included",
            "Included",
        ],
        "HackerNews": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
        "Ubuntu IRC": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
        "EuroParl": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
        ],
        "StackExchange": [
            "Included",
            "-",
            "-",
            "-",
            "-",
            "-",
            "Included",
            "Included",
        ],
        "Code": [
            "- what is this?",
            "-",
            "-",
            "-",
            "-",
            "Included",
            "Included",
            "Included",
        ],
    }
)
# Apply table styling: Light green for the header, alternating white and light grey for rows
styled_table = (
    dataset_comparison2.style.set_properties(
        **{"background-color": "#E1EEDB"},
        subset=pd.IndexSlice[0, :],  # Row 0 with a light green background
    )
    .apply(
        lambda x: [
            "background-color: #E1EEDB"
            if i == 0
            else (
                "background-color: rgb(237, 242, 251)"
                if i % 2 == 0
                else "background-color: white"
            )
            for i in range(len(x))
        ],
        axis=0,
    )
    .hide(axis="index")
)  # Hide the row index

# Use _repr_html_() method to get the HTML representation of the styled DataFrame
table_html2 = styled_table._repr_html_()
# table_html2 = dataset_comparison2.to_html(index=False, border=0)
table_div_2 = Div(NotStr(table_html2), style="margin: 40px;")

dataset_sources = pd.DataFrame(
    {
        "Data Source": [
            "CommonCrawl",
            "Papers",
            "Wikipedia",
            "Freelaw",
            "DM Math",
            "USPTO",
            "PG-19",
            "HackerNews",
            "Ubuntu IRC",
            "Europarl",
            "StackExchange",
        ],
        "Raw Data Size": [
            "11 TB",
            "712 GB",
            "210 GB",
            "23 GB",
            "22 GB",
            "45 GB",
            "11 GB",
            "4.1 GB",
            "4.7 GB",
            "6.1 GB",
            "45 GB",
        ],
        "Token Count": [
            "5.71T",
            "154.96B",
            "4.75B",
            "7.34B",
            "5.23B",
            "4.95B",
            "2.94B",
            "1.08B",
            "1.54B",
            "1.96B",
            "8.37B",
        ],
        "Cut-Off Date": [
            "2024-30",
            "Q4 2023",
            "-",
            "Q1 2024",
            "-",
            "Q4 2023",
            "-",
            "Q4 2023",
            "Q4 2023",
            "-",
            "Q4 2023",
        ],
    }
)
# Apply table styling: Light green for the header, alternating white and light grey for rows
styled_table = dataset_sources.style.apply(
    lambda x: [
        "background-color: white"
        if i % 2 == 0
        else "background-color: rgb(237, 242, 251)"
        for i in range(len(x))
    ],
    axis=0,
).hide(axis="index")  # Hide the row index

table_html_data = styled_table._repr_html_()
# table_html_data = dataset_sources.to_html(index=False, border=0)
table_div_data = Div(NotStr(table_html_data), style="margin: 40px;")


@app.get("/intro")
def intro():
    return Div(
        Section(
            H2("About TxT360"),
            P(
                B("We introduce TxT360 (Trillion eXtracted Text) the first dataset to globally deduplicate 99 CommonCrawl snapshots and 14 commonly used non-web data sources (e.g. FreeLaw, PG-19, etc.) providing pretraining teams with a recipe to easily adjust data weighting and train the most performant models.")
            ),
            P(
                "Building on top of the prior studies on pre-training data, TxT360 carefully implements data processing steps including extraction, filtering, deduplication, personally identifiable information removal, and other steps."
            ),
            P(
                "Metadata is stored to recover the raw distribution for each dataset, enabling fine-grained control to create data distributions and corpus of desired size. As an example, we present one simple upsampling scheme that takes into account the duplication counts, resulting in a 15~16 trillion token corpus, outperforming FineWeb and our non-upsampling baselines, on diverse evaluations. Unlike DCLM and RedPajama V2, we present the final deduplicated dataset that is ready to go."
            ),
            P(
                "We documented all implementation details in this blog post and are open sourcing the code. Examples of each filter and rationale supporting each decision are included."
            ),
            id="section1",
        ),
        Section(
            H2("Motivation Behind Txt360"),
            H3(
                "TxT360 is the first dataset to combine both web and curated data sources commonly used in pretraining."
            ),
            table_div_1,
            table_div_2,
            P(
                "In pretraining, it is common to combine web data and curated sources (cite). Web data is included to provide a vast quantity of long tail and diverse data, while curated datasets are often information rich and provide the 'deep-dive' domain information. Combining both datasets plays a critical role for effective LLM pre-training. By integrating the reach of web data with the quality of curated sources, TxT360 meets and surpasses the rigorous standards required for state-of-the-art LLM pre-training. See Results section below."
            ),
            #P("Table 2: Basic TxT360 Statistics."),
            #table_div_data,
            id="section2",
        ),
        Section(
            H2("Our Generalizable Approach to Data Processing"),
            P(
                "To produce TxT360, a comprehensive and transparent data processing pipeline was designed to account for the nuances of both web and curated datasets. The pipeline presents a unified framework for processing both data types, making it convenient and easily adaptive for users to revise and fine-tune the pipeline for their own use cases."
            ),
            P(
                "Web datasets are inherently noisy and varied. The TxT360 pipeline implements sophisticated filtering and deduplication techniques to clean and remove redundancies while preserving data integrity."
            ),
            P(
                "Curated datasets are typically structured and consistently formatted. TxT360 filters these sources with selective steps to maintain their integrity while providing seamless integration into the larger dataset. Both data source types are globally deduplicated together resulting in 5.7T tokens of high-quality data."
            ),
            P(
                "We provide details and context for the choices behind TxT360 in the respective Web Data Processing and Curated Source Processing section. A deep dive in the deduplication [here]. "
            ),
            #Img(src="images/pipeline.png", height="300", width="600"),
            #P(
            #    "Figure 1: Data processing pipeline. All the steps are adopted for processing web data while the yellow blocks are adopted for processing curated sources."
            #),
            id="section3",
        ),
        id="inner-text",
    )


rt("/update/{target}")(data_viewer.update)

rt("/curated")(curated.curated)

rt("/webdata")(web.web_data)

rt("/common")(common.common_steps)

rt("/results")(results.results)

serve()