pipeline/config.py

from typing import List

# this try/except is important for publishing to Hugging Face
try:
    from dagster import Config
except ImportError:
    Config = object


class LLMBoardConfig(Config):
    group_columns: List[str] = ["model", "language", "template_name"]
    single_values_columns: List[str] = ["execution_time", "characters_count", "words_count"]
    list_columns: List[str] = ["chunk_sizes", "chunk_generation_times", "chunk_generation_times_by_chunk_sizes"]
    plot_dir: str = "./html/plots/"
    plot_json_dir: str = "./data/"
    saving_path: str = "data/"


class QueriesConfig(Config):
    base_query_template: str = """Summarize me this text, the summary should be in {language}
```
{text}
```
"""

    query_template: dict = {
        "markdown": """Return output as markdown""",
        "json": """Return output as json in format:
{
"summary": "<summary">
}""",
        "call": """Return output by calling summary_result()""",
    }


class MeasurementsConfig(Config):
    mock: bool = False
    remove_old_measurements: bool = False
    small_dataset: bool = False


class QueriesDatasetConfig(Config):
    dataset_name: str = "GEM/xlsum"
    samples_per_measurement: int = 20
    languages: List[str] = ["english", "japanese"]
    query_config: QueriesConfig = QueriesConfig()


class SummaryConfig(Config):
    saving_path: str = "data/"


class ModelCostsConfig(Config):
    saving_path: str = "data/"


class TimeOfDayComparisonConfig(Config):
    saving_path: str = "data/"


class GeneralPlotConfig(Config):
    plots_dir: str = "./html/plots/"
    saving_path: str = "data/"
    seconds_per_token: float = 184 / 6
    input_size: int = 100
    expected_output_size: int = 50