app.py

import os
import json
import datetime
import smtplib
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.utils import parseaddr
import numpy as np
import gradio as gr
import pandas as pd
from datasets import load_dataset
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import HfApi


TOKEN = os.environ.get("TOKEN", None)
SUBMISSION_DATASET = "KoLMogorov-Test/submissions"
VERSION = "v1"

api = HfApi()

        
def format_error(msg):
    return f"<p style='color: red; font-size: 20px; text-align: center;'>{msg}</p>"

def format_warning(msg):
    return f"<p style='color: orange; font-size: 20px; text-align: center;'>{msg}</p>"

def format_log(msg):
    return f"<p style='color: green; font-size: 20px; text-align: center;'>{msg}</p>"

def model_hyperlink(link, model_name):
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{model_name}</a>'


# Function to restart the space
def restart_space():
    return

TYPES = ["markdown", "markdown", "number", "number", "number", "number", "number", "number", "str", "str"]


def save_files_for_eval(
    model_name: str,
    organization: str,
    organization_email: str,
    path_to_file,
    prior_file,
    results_json,
    modality: str,
    split: str,
    token: str,
):
    # Validate inputs
    if not model_name:
        return format_error("Model name is required.")
    if not organization and organization_email:
        return format_error("Organization and email  required.")
    
    # Get the current date for the directory name
    date_str = datetime.datetime.today().strftime('%Y-%m-%d')
    
    # Define the base path in the repository (remove leading slash)
    base_path_in_repo = f"{organization}/{model_name}/{VERSION}_{split}_{modality}_{date_str}"
    
    # List of files to upload
    files_to_upload = [
        (path_to_file, "prediction.jsonl"),
        (prior_file, "decoder.py"),
        (results_json, "results.json")
    ]
    
    # Upload each file
    for file, filename in files_to_upload:
        if file is not None:
            path_in_repo = f"{base_path_in_repo}/{filename}"
            api.upload_file(
                repo_id=SUBMISSION_DATASET,
                path_or_fileobj=file.name,
                path_in_repo=path_in_repo,
                repo_type="dataset",
                token=token
            )
        else:
            return format_error(f"File {filename} is missing.")

    return format_log("All files uploaded successfully. Please email us at [email protected] to verify your submission.")


# Gradio interface
demo = gr.Blocks()
with demo:
    gr.HTML("""<h1 id="space-title">The KoLMogorov-Test: Can CodeLMs lead to the next breakthrough in data compression?</h1>""")
    gr.HTML("""The Kolmogorov complexity of a sequence is the length of the shortest computer program that produces the sequence. <br><br>The aim of the KoLMogorov-Test (KT) is to empirically evaluate the ability of CodeLMs to detect patterns in and compress sequences by writing short programs that output them.
<h2>The Task</h2>
Given a sequence, the CodeLM is asked to produce a short python program that outputs the sequence. The programs are compressed by a user-defined compressor before submission, and code length is measured in compressed form. In order to evaluate the correctness of a program, it is first decoded using a user-provided decoder and then executed. The decoder size counts towards the compressed length, and may contain e.g. a library of helper functions (DSL). <br><br>

The length of python and its standard library are not counted, but in order to prevent cheating the programs must pass a stringent check. Programs cannot import arbitrary modules, use advanced language features, access the internet, etc. <br><br>

KT currently includes six modalities - text, DNA, three encodings of audio data (MFCC, 16-bit, and 8-bit), and synthetic sequences produced by random programs. Two dataset sizes are available: a small one with 1MB per modality, and a large one with 1GB (DNA and text only).""")

    gr.HTML("""<h2>Getting Started</h2>
            Access the data from the <a href="https://github.com/facebookresearch/KoLMogorov" target="_blank">GitHub repo</a>. 
            <br><br>In addition, we provide code to reproduce experiments from the paper including our DSL and evaluation code.
            <br><br>If you have any questions, please email us at <a href="mailto:[email protected]">[email protected]</a>.
    <br><br>You are allowed to:<br><br>
    <ul>
        <li>Split the original sequence to subsequent sub-sequences.</li>
        <li>Use the standard python library including Gzip.</li>
        <li>Create new DSLs.</li>
        <li>Create new priors to encode the programs.</li>
    </ul>

    <br>With the following restrictions:<br><br>
    <ul>
        <li>To prevent cases where compression of the data is performed by external code, we prevent access to the internet and usage of previous compression algorithms, excluding Gzip which is viewed as a strong baseline. </li>
        <li>We return the compression rate with and without the cost of the decoder. When reporting results, we consider the additional code as negligible if it is &lt100KB and does not scale with the length of the sequence. When reporting results for the 1GB seqeunces, please always report results including the decoder.</li>
    </ul>
""")

    gr.HTML("""<h2>How can we achieve future progress?</h2>""")
    gr.Markdown("""In the paper we show that stronger models perform better on KT and thatn CodeLMs can outperform other compression methods on synthetic distributions when training data is avialable. Another exciting direction for future work is to encode programs as lambda experssions, making it feasible to fit a small interpreter of only 383 in the decoder as <a href="https://justine.lol/lambda/" target="_blank">suggested here</a>.""")
    
    gr.HTML("""<h2>Making a New Submission</h2>""")
            
    gr.Markdown("""To make a new submission, upload three files - a file with the encoded program, a file that decodes the programs to executable code, and a file with the expected results. Let's consider a toy example for the sequence - [5, 10, 13, 14, 16, 5, 5, 5]""")
    
    gr.Markdown("""The programs file is a jsonl that follows the following format. We recommend submitting a single program for the whole sequence:
            ```
            {"sub_sequence_start_index": 0, "sub_sequence_end_index": 7 (the length of the sequence), "encoded_program": "H4sIAI7Te2cC/8svLSkoLVGwVYg21VEwNABiYyA2AWIzHQVTCIoFAMxnYTIlAAAA (An encoding of the program whose execution results in the input sequence. This is the Gzip encoding of the program that returns the input sequence)."}
            ```

            For simplicity, we allow splitting to subsequent sub-sequences:
            ```
            {"sub_sequence_start_index": 0, "sub_sequence_end_index": k, "encoded_program": "An encoding of the program whose execution results in the subsequence between indeses [0, n]."}
            {"sub_sequence_start_index": k+1, "sub_sequence_end_index": k+1+j, "encoded_program": "An encoding of the program whose execution results in the subsequence between indeses [n+1, n+1+m]."}
            ...
            {"sub_sequence_start_index": n, "sub_sequence_end_index": len(sequence), "encoded_program": "An encoding of the program whose execution results in the subsequence between indeses [n, len(sequence)]."}
            ```
        """)

    gr.Markdown("""For decoding, upload a python file that implements the 'decode' method, which receives as input an programs from the programs file and returns the executable python program. Executing the decoded program must result in the input sequence. For example, this is a decoder that decompresses using Gzip. 
            ```
            import gzip
            import base64

            def decode(program):
                return gzip.decompress(base64.b64decode(program)).decode('utf-8')
            ```
        """)
    
    gr.Markdown("""For the results, upload a Json file that matches the Result object returned from the <a href="https://github.com/facebookresearch/KoLMogorov/tree/main/src/experiments/evaluation" target="_blank">official evaluation script</a>. If you are interested in verification of your results, please send an <a href="mailto:[email protected]">email</a> with the submission details. We will then execute the code and verify the execution matches the original sequece.
            ```
            {
                "compressed_programs_size": 64,
                "decoder_size": 118,
                "compressed_size": 182,
                "compression_rate_without_decoder": 8.0,
                "compression_rate": 22.75,
                "accuracy": 1,
                "gold_data_size": 8,
                "first_error": null
            }
            ```
        """)
    
    with gr.Accordion(""):
        with gr.Row():
            with gr.Column():
                model_name_textbox = gr.Textbox(label="Model Name")
                organization = gr.Textbox(label="Organization")
                mail = gr.Textbox(
                    label="Contact Email (will be stored privately)"
                )
                split_dropdown = gr.Dropdown(
                    choices=["Small (1MB)", "Large (1GB)"],
                    label="Size",
                )
                modality_dropdown = gr.Dropdown(
                    choices=["Text", "Dna", "Audio-MFCC", "Audio-16-Bit", "Audio-8-Bit", "Synthetic"],
                    label="Modality",
                )
                prior_dropdown = gr.Dropdown(
                    choices=["Custom (attached file)", "None"],
                    label="Decoding",
                )
                url_textbox = gr.Textbox(label="URL to Project Information",
                                         interactive=True,)
                model_family_textbox = gr.Textbox(label="Base Model")

            with gr.Column():
                file_output = gr.File(label="Programs File")
                prior_file = gr.File(label="Decoder File")
                results_json = gr.File(label="Results Json")
                
        submit_button = gr.Button("Submit")
        submission_result = gr.Markdown()

    submit_button.click(
        lambda model_name, split, model_family, url, path_to_file, organization, mail, prior_file, results_json, prior_dropdown, modality_dropdown: (
            save_files_for_eval(
                model_name,
                organization,
                mail,
                path_to_file,
                prior_file,
                results_json,
                modality_dropdown,
                split,
                TOKEN  # Ensure TOKEN is defined and accessible
            )
        ),
        [
            model_name_textbox,
            split_dropdown,
            model_family_textbox,
            url_textbox,
            file_output,
            organization,
            mail,
            prior_file,
            results_json,
            prior_dropdown,
            modality_dropdown
        ],
        submission_result,
    )

    data = {
        "Text": {
            "Model": ["Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 69.5, 33.3, 18.0, 8.5],
            "Precision": [1.0, 1.34, 2.18, 2.78, 2.48],
            "Compression Rate": [0.357, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓'],
        },
        "DNA": {
            "Model": ["Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 54.2, 6.5, 9.6, 1.4],
            "Precision": [1.0, 1.94, 3.17, 3.17, 3.12],
            "Compression Rate": [0.714, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓'],
        },
        "Audio-8-bit": {
            "Model": ["Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 36.4, 15.0, 10.1, 3.9],
            "Precision": [1.0, 1.43, 1.66, 1.67, 1.74],
            "Compression Rate": [0.398, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓'],
        },
        "Audio-16-bit": {
            "Model": ["Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 69.5, 35.6, 18.0, 5.9],
            "Precision": [1.0, 1.34, 1.66, 1.96, 1.54],
            "Compression Rate": [0.920, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓'],
        },
        "Audio-MFCC": {
            "Model": ["Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 83.8, 29.6, 24.2, 8.8],
            "Precision": [1.0, 1.33, 1.58, 1.56, 1.51],
            "Compression Rate": [0.903, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓'],
        },
        "Synthetic": {
            "Model": ["SeqCoder-8B + Gzip", "Gzip", "GPT4-o", "LLAMA-3.1-405B", "LLAMA-3.1-70B", "LLAMA-3.1-8B"],
            "Pass@1": [100, 100, 44.7, 24.8, 22.5, 3.7],
            "Precision": [0.64, 1.0, 1.65, 2.06, 2.18, 2.34],
            "Compression Rate": [0.38, 0.593, "n/a", "n/a", "n/a", "n/a"],
            "Verified": ['✓', '✓', '✓', '✓', '✓', '✓'],
        },
    }
    
    def refresh():
        # Add refresh logic
        return 
    
    gr.HTML("<h2>KT Leaderboard - Can you beat Gzip on KT?</h2>")
    for k,v in data.items():
        with gr.Tab(k):
            leaderboard_table_test = gr.Dataframe(
                value=pd.DataFrame(data[k]),
                interactive=False,
            )

    refresh_button = gr.Button("Refresh")
    refresh_button.click(
        refresh,
        inputs=[],
        outputs=[leaderboard_table_test],
    )
     
    with gr.Row():
        with gr.Accordion("📙 Citation", open=False):
            citation_text = """@inproceedings{
anonymous2024the,
title={The Ko{LM}ogorov Test: Compression by Code Generation},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=C45YqeBDUM},
note={under review}
}"""
            citation_button = gr.Textbox(
                value=citation_text,
                label="Citation",
                lines=20,
                elem_id="citation-button",
                show_copy_button=True
            )

    gr.HTML(
        "<p>We would like to thank the GAIA team for sharing the source code for their leaderboard which we used as a template and HuggingFace for hosting the leaderboard.</p>")

scheduler = BackgroundScheduler()
scheduler.add_job(restart_space, "interval", seconds=3600)
scheduler.start()
demo.launch(debug=True)