autoddg 0.1.0__py3-none-any.whl

autoddg/__init__.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from .autoddg import AutoDDG
+from .description import DatasetDescriptionGenerator, SearchFocusedDescription
+from .evaluation import (
+    GPT4oMiniPairwiseEvaluator,
+    GPT4oPairwiseEvaluator,
+    GPTEvaluator,
+    LLaMAEvaluator,
+    LLaMAPairwiseEvaluator,
+    PairwiseEvaluator,
+)
+from .llm import LLMClient, LocalLLMClient, OpenAICompatibleClient
+from .profiling import SemanticProfiler, profile_dataset
+from .topic import DatasetTopicGenerator
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AutoDDG",
+    "DatasetDescriptionGenerator",
+    "DatasetTopicGenerator",
+    "GPTEvaluator",
+    "LLaMAEvaluator",
+    "PairwiseEvaluator",
+    "GPT4oPairwiseEvaluator",
+    "GPT4oMiniPairwiseEvaluator",
+    "LLaMAPairwiseEvaluator",
+    "LLMClient",
+    "LocalLLMClient",
+    "OpenAICompatibleClient",
+    "SearchFocusedDescription",
+    "SemanticProfiler",
+    "profile_dataset",
+    "__version__",
+]

autoddg/autoddg.py

@@ -0,0 +1,324 @@
+from __future__ import annotations
+
+from typing import Any
+
+from beartype import beartype
+from pandas import DataFrame
+
+from .description import DatasetDescriptionGenerator, SearchFocusedDescription
+from .evaluation import BaseEvaluator
+from .llm import LocalLLMClient, OpenAICompatibleClient
+from .profiling import SemanticProfiler, profile_dataset
+from .topic import DatasetTopicGenerator
+
+
+@beartype
+class AutoDDG:
+    """AutoDDG - Automated Dataset Description Generator
+
+    AutoDDG is the library's entry class, exposing:
+    * profiling,
+    * semantic analysis,
+    * topic generation,
+    * dataset description,
+    * search-focused description expansion, and
+    * optional description evaluation
+
+    Args:
+        client (Any): OpenAI-compatible client (e.g. ``openai.OpenAI(...)``) or None
+            for local LLM.
+        model_name (str): Model identifier (e.g. ``"gpt-4o"`` for API or
+            ``"Qwen/Qwen2.5-7B-Instruct"`` for local).
+        use_local_llm (bool): If True, use local LLM via transformers.
+            Requires transformers and torch.
+        local_llm_device (str | None): Device for local LLM ("cuda", "cpu", or None
+            for auto).
+        local_llm_dtype (str | None): Data type for local LLM ("float16", "bfloat16",
+            "float32", or None for auto).
+        description_temperature (float): Temperature for description generation.
+        description_words (int): Target word count for generated descriptions.
+        search_model_name (str | None): Override model for search-expansion.
+        semantic_model_name (str | None): Override model for semantic profiling.
+        topic_temperature (float): Temperature for topic generation.
+        evaluator (BaseEvaluator | None): Optional evaluator for quality scoring.
+
+    Examples:
+        Basic usage with API:
+
+            >>> import openai
+            >>> from autoddg import AutoDDG
+            >>> client = openai.OpenAI(api_key="sk-...")
+            >>> pipe = AutoDDG(
+            ...     client=client, model_name="gpt-4o", description_words=100
+            ... )
+            >>> sample_csv = (
+            ...     "city,country,population\\nLondon,UK,8908081\\nLeeds,UK,789194"
+            ... )
+            >>> prompt, desc = pipe.describe_dataset(dataset_sample=sample_csv)
+            >>> print(desc)
+
+        Basic usage with local LLM:
+
+            >>> from autoddg import AutoDDG
+            >>> pipe = AutoDDG(
+            ...     client=None,
+            ...     model_name="Qwen/Qwen2.5-7B-Instruct",
+            ...     use_local_llm=True,
+            ...     description_words=100
+            ... )
+            >>> sample_csv = (
+            ...     "city,country,population\\nLondon,UK,8908081\\nLeeds,UK,789194"
+            ... )
+            >>> prompt, desc = pipe.describe_dataset(dataset_sample=sample_csv)
+            >>> print(desc)
+
+        Advanced usage with topic and evaluator:
+
+            >>> import pandas as pd
+            >>> from autoddg import GPTEvaluator
+            >>> df = pd.DataFrame({
+            ...     "city": ["London", "Leeds"],
+            ...     "country": ["UK", "UK"],
+            ...     "population": [8908081, 789194],
+            ... })
+            >>> profile, semantic = pipe.profile_dataframe(df)
+            >>> topic = pipe.generate_topic("UK Cities", None, df.to_csv(index=False))
+            >>> _, desc = pipe.describe_dataset(
+            ...     dataset_sample=df.to_csv(index=False),
+            ...     dataset_profile=profile,
+            ...     use_profile=True,
+            ...     semantic_profile=semantic,
+            ...     use_semantic_profile=True,
+            ...     data_topic=topic, use_topic=True,
+            ... )
+            >>> evaluator = GPTEvaluator(gpt4_api_key="sk-...")
+            >>> pipe.set_evaluator(evaluator)
+            >>> scores = pipe.evaluate_description(desc)
+            >>> print(scores)
+    """
+
+    def __init__(
+        self,
+        client: Any | None = None,
+        model_name: str = "gpt-4o",
+        *,
+        use_local_llm: bool = False,
+        local_llm_device: str | None = None,
+        local_llm_dtype: str | None = None,
+        description_temperature: float = 0.0,
+        description_words: int = 100,
+        search_model_name: str | None = None,
+        semantic_model_name: str | None = None,
+        topic_temperature: float = 0.0,
+        evaluator: BaseEvaluator | None = None,
+    ) -> None:
+        # Initialize LLM client
+        if use_local_llm:
+            if client is not None:
+                raise ValueError(
+                    "Cannot specify both client and use_local_llm=True. "
+                    "For local LLM, set client=None and use_local_llm=True."
+                )
+            llm_client = LocalLLMClient(
+                model_name=model_name,
+                device=local_llm_device,
+                torch_dtype=local_llm_dtype,
+            )
+        elif client is not None:
+            llm_client = OpenAICompatibleClient(client)
+        else:
+            raise ValueError(
+                "Must provide either client (for API) or set use_local_llm=True " "(for local LLM)"
+            )
+
+        self.client = client  # Keep for backward compatibility
+        self.model_name = model_name
+        self.description_generator = DatasetDescriptionGenerator(
+            client=llm_client,
+            model_name=model_name,
+            temperature=description_temperature,
+            description_words=description_words,
+        )
+        self.semantic_profiler = SemanticProfiler(
+            client=llm_client,
+            model_name=semantic_model_name or model_name,
+        )
+        self.topic_generator = DatasetTopicGenerator(
+            client=llm_client,
+            model_name=model_name,
+            temperature=topic_temperature,
+        )
+        self.search_description = SearchFocusedDescription(
+            client=llm_client,
+            model_name=search_model_name or model_name,
+        )
+        self.evaluator = evaluator
+
+    def describe_dataset(
+        self,
+        dataset_sample: str,
+        dataset_profile: str | None = None,
+        use_profile: bool = False,
+        semantic_profile: str | None = None,
+        use_semantic_profile: bool = False,
+        data_topic: str | None = None,
+        use_topic: bool = False,
+    ) -> tuple[str, str]:
+        """
+        Produce a short description from a CSV sample with optional context
+
+        Args:
+            dataset_sample: CSV text containing example rows
+            dataset_profile: Structural profile text
+            use_profile: Include the structural profile if True
+            semantic_profile: Natural-language column semantics
+            use_semantic_profile: Include the semantic profile if True
+            data_topic: Short topic string for the dataset
+            use_topic: Include the topic if True
+
+        Returns:
+            (prompt, description)
+        """
+
+        return self.description_generator.generate_description(
+            dataset_sample=dataset_sample,
+            dataset_profile=dataset_profile,
+            use_profile=use_profile,
+            semantic_profile=semantic_profile,
+            use_semantic_profile=use_semantic_profile,
+            data_topic=data_topic,
+            use_topic=use_topic,
+        )
+
+    def profile_dataframe(self, dataframe: DataFrame) -> tuple[str, str]:
+        """
+        Summarise structure and coverage using the datamart profiler
+
+        Ref: https://pypi.org/project/datamart-profiler/
+
+        Args:
+            dataframe: Input frame
+
+        Returns:
+            (profile_text, semantic_notes)
+        """
+
+        return profile_dataset(dataframe)
+
+    def analyze_semantics(
+        self,
+        dataframe: DataFrame,
+        *,
+        use_group_prompting: bool = False,
+        use_multi_threading: bool = False,
+        use_batch_processing: bool = False,
+        max_workers: int | None = None,
+        group_size: int = 0,
+        batch_size: int = 4,
+    ) -> str:
+        """
+        Infer column semantics with an LLM and return a short overview
+
+        Processing modes available:
+        1. Sequential mode (default): Processes columns one by one sequentially.
+        2. Multi-threaded mode (use_multi_threading=True): Uses multi-threading to
+           process columns in parallel for faster execution. Only available for
+           OpenAI API clients.
+        3. Group mode (use_group_prompting=True): Processes columns in groups via
+           group prompting, reducing API calls.
+           - If group_size=0: Processes all columns in a single prompt (most
+             efficient).
+           - If group_size>0: Processes columns in groups of group_size.
+        4. Batch mode (use_batch_processing=True): Processes columns in batches
+           using batch inference. Only available for local LLMs. Takes precedence
+           over other modes.
+
+        Args:
+            dataframe: Input frame
+            use_group_prompting: If True, use group prompting (single API call for all
+                columns or groups). Takes precedence over use_multi_threading.
+            use_multi_threading: If True, use multi-threading for individual column
+                processing (only used if use_group_prompting=False and
+                use_batch_processing=False). Only works with OpenAI API clients.
+            use_batch_processing: If True, use batch processing for local LLMs.
+                Only works with LocalLLMClient. Takes precedence over other modes.
+            max_workers: Maximum number of concurrent workers for multi-threaded mode.
+                Default: min(32, num_columns).
+            group_size: Group size for group prompting. If 0, process all columns
+                at once. If >0, process in groups of that size.
+            batch_size: Batch size for batch processing mode. Only used with local LLMs.
+
+        Returns:
+            Summary of column semantics
+        """
+
+        return self.semantic_profiler.analyze_dataframe(
+            dataframe,
+            use_group_prompting=use_group_prompting,
+            use_multi_threading=use_multi_threading,
+            use_batch_processing=use_batch_processing,
+            max_workers=max_workers,
+            group_size=group_size,
+            batch_size=batch_size,
+        )
+
+    def generate_topic(
+        self, title: str, original_description: str | None, dataset_sample: str
+    ) -> str:
+        """
+        Generate a 2–3 word topic from title description and sample
+
+        Args:
+            title: Dataset title
+            original_description: Existing description if available
+            dataset_sample: CSV text sample
+
+        Returns:
+            Short topic string
+        """
+
+        return self.topic_generator.generate_topic(title, original_description, dataset_sample)
+
+    def expand_description_for_search(self, description: str, topic: str) -> tuple[str, str]:
+        """
+        Expand a readable description into a search-oriented variant
+
+        Args:
+            description: Original dataset description
+            topic: Topic string
+
+        Returns:
+            (prompt, expanded_description)
+        """
+
+        return self.search_description.expand_description(description, topic)
+
+    def evaluate_description(self, description: str) -> str:
+        """
+        Score a description with the configured evaluator
+
+        Args:
+            description: Description text to score
+
+        Returns:
+            Evaluation response
+
+        Raises:
+            RuntimeError: If no evaluator is set
+        """
+
+        if self.evaluator is None:
+            raise RuntimeError(
+                "No evaluator configured for AutoDDG. Provide one via set_evaluator()."
+            )
+        return self.evaluator.evaluate(description)
+
+    def set_evaluator(self, evaluator: BaseEvaluator) -> None:
+        """
+        Attach or replace the evaluator to use for scoring
+
+        Args:
+            evaluator: Evaluator instance
+        """
+
+        self.evaluator = evaluator

autoddg/configurations/prompts.yaml

@@ -0,0 +1,236 @@
+dataset_description:
+  introduction: |
+    Answer the question using the following information.
+
+    First, consider the dataset sample:
+
+    {dataset_sample}
+  profile_instruction: |
+    Additionally, the dataset profile is as follows:
+
+    {dataset_profile}
+
+    Based on this profile, please add sentence(s) to enrich the dataset description.
+  semantic_instruction: |
+    Furthermore, the semantic profile of the dataset columns is as follows:
+    {semantic_profile}
+
+    Based on this information, please add sentence(s) discussing the semantic profile in the description.
+  topic_instruction: |
+    Moreover, the dataset topic is: {data_topic}. Based on this topic, please add sentence(s) describing what this dataset can be used for.
+  closing_instruction: |
+    Question: Based on the information above and the requirements, provide a dataset description in sentences. Use only natural, readable sentences without special formatting.
+
+    Answer:
+  system_message: |
+    You are an assistant for a dataset search engine. Your goal is to improve the readability of dataset descriptions for dataset search engine users.
+
+search_expansion:
+  template: |
+    Dataset Overview:
+    - Please keep the exact initial description of the dataset as shown in beginning the prompt.
+
+    Key Themes or Topics:
+    - Central focus on a broad area of interest (e.g., urban planning, socio-economic factors, environmental analysis).
+    - Data spans multiple subtopics or related areas that contribute to a holistic understanding of the primary theme.
+    Example:
+    - theme1/topic1
+    - theme2/topic2
+    - theme3/topic3
+
+    Applications and Use Cases:
+    - Facilitates analysis for professionals, policymakers, researchers, or stakeholders.
+    - Useful for specific applications, such as planning, engineering, policy formulation, or statistical modeling.
+    - Enables insights into patterns, trends, and relationships relevant to the domain.
+    Example:
+    - application1/usecase1
+    - application2/usecase2
+    - application3/usecase3
+
+    Concepts and Synonyms:
+    - Includes related concepts, terms, and variations to ensure comprehensive coverage of the topic.
+    - Synonyms and alternative phrases improve searchability and retrieval effectiveness.
+    Example:
+    - concept1/synonym1
+    - concept2/synonym2
+    - concept3/synonym3
+
+    Keywords and Themes:
+    - Lists relevant keywords and themes for indexing, categorization, and enhancing discoverability.
+    - Keywords reflect the dataset's content, scope, and relevance to the domain.
+    Example:
+    - keyword1
+    - keyword2
+    - keyword3
+
+    Additional Context:
+    - Highlights the dataset's relevance to specific challenges or questions in the domain.
+    - May emphasize its value for interdisciplinary applications or integration with related datasets.
+    Example:
+    - context1
+    - context2
+    - context3
+  user_prompt: |
+    You are given a dataset about the topic {topic}, with the following initial description:
+
+    {initial_description}
+
+    Please expand the description by including the exact topic. Additionally, add as many related concepts, synonyms, and relevant terms as possible based on the initial description and the topic.
+
+    Unlike the initial description, which is focused on presentation and readability, the expanded description is intended to be indexed at the backend of a dataset search engine to improve searchability.
+
+    Therefore, focus less on readability and more on including all relevant terms related to the topic. Make sure to include any variations of the key terms and concepts that could help improve retrieval in search results.
+
+    Please follow the structure of the following example template:
+    {template}
+  system_message: |
+    You are an assistant for a dataset search engine. Your goal is to improve the performance of the dataset search engine for keyword queries.
+
+semantic_profiler:
+  template: |
+    {'Temporal':
+        {
+            'isTemporal': Does this column contain temporal information? Yes or No,
+            'resolution': If Yes, specify the resolution (Year, Month, Day, Hour, etc.).
+        },
+     'Spatial': {'isSpatial': Does this column contain spatial information? Yes or No,
+                 'resolution': If Yes, specify the resolution (Country, State, City, Coordinates, etc.).},
+     'Entity Type': What kind of entity does the column describe? (e.g., Person, Location, Organization, Product),
+     'Domain-Specific Types': What domain is this column from (e.g., Financial, Healthcare, E-commerce, Climate, Demographic),
+     'Function/Usage Context': How might the data be used (e.g., Aggregation Key, Ranking/Scoring, Interaction Data, Measurement).}
+  response_example: |
+    {
+    "Domain-Specific Types": "General",
+    "Entity Type": "Temporal Entity",
+    "Function/Usage Context": "Aggregation Key",
+    "Spatial": {"isSpatial": false,
+                "resolution": ""},
+    "Temporal": {"isTemporal": true,
+                "resolution": "Year"}
+    }
+  user_prompt: |
+    You are a dataset semantic analyzer. Based on the column name and sample values, classify the column into multiple semantic types.
+    Please group the semantic types under the following categories:
+    'Temporal', 'Spatial', 'Entity Type', 'Data Format', 'Domain-Specific Types', 'Function/Usage Context'.
+    Following is the template {template}
+    Please follow these rules:
+    1. The output must be a valid JSON object that can be directly loaded by json.loads. Example response is {response_example}
+    2. All keys from the template must be present in the response.
+    3. All keys and string values must be enclosed in double quotes.
+    4. There must be no trailing commas.
+    5. Use booleans (true/false) and numbers without quotes.
+    6. Do not include any additional information or context in the response.
+    7. If you are unsure about a specific category, you can leave it as an empty string.
+
+    Column name: {column_name}
+    Sample values: {sample_values}
+  system_message: |
+    You are a helpful assistant skilled in dataset semantic analysis.
+  group_response_example: |
+    {
+      "column1": {
+        "Domain-Specific Types": "General",
+        "Entity Type": "Temporal Entity",
+        "Function/Usage Context": "Aggregation Key",
+        "Spatial": {"isSpatial": false, "resolution": ""},
+        "Temporal": {"isTemporal": true, "resolution": "Year"}
+      },
+      "column2": {
+        "Domain-Specific Types": "Financial",
+        "Entity Type": "Numeric Entity",
+        "Function/Usage Context": "Measurement",
+        "Spatial": {"isSpatial": false, "resolution": ""},
+        "Temporal": {"isTemporal": false, "resolution": ""}
+      }
+    }
+  group_user_prompt: |
+    You are a dataset semantic analyzer. Analyze multiple columns at once and classify each column into multiple semantic types.
+
+    For each column provided below, classify it into semantic types under the following categories:
+    'Temporal', 'Spatial', 'Entity Type', 'Data Format', 'Domain-Specific Types', 'Function/Usage Context'.
+
+    Template for each column:
+    {template}
+
+    Please follow these rules:
+    1. The output must be a valid JSON object where each key is a column name and each value is the semantic type classification for that column.
+    2. All keys from the template must be present in each column's classification.
+    3. All keys and string values must be enclosed in double quotes.
+    4. There must be no trailing commas.
+    5. Use booleans (true/false) and numbers without quotes.
+    6. Do not include any additional information or context in the response.
+    7. If you are unsure about a specific category for a column, you can leave it as an empty string.
+
+    Example response format:
+    {response_example}
+
+    Columns to analyze:
+    {columns_text}
+
+    Return a JSON object with column names as keys and their semantic classifications as values:
+
+topic_generation:
+  system_message: |
+    You are an assistant for generating concise dataset topics.
+  user_prompt: |
+    Using the dataset information provided, generate a concise topic in 2-3 words that best describes the dataset's primary theme.
+
+    Title: {title}
+    {description_block}
+    Dataset Sample: {dataset_sample}
+
+    Topic (2-3 words):
+
+evaluation:
+  system_message: |
+    You are a helpful and precise assistant for checking the quality of the dataset description.
+  user_prompt: |
+    You will be given one tabular dataset description. Your task is to rate the description on 3 metrics.
+    Please make sure you read and understand these instructions carefully. Please keep this document open while reviewing, and refer to it as needed.
+
+    Evaluation Criteria:
+    1. Completeness (1-10) - Evaluates how thoroughly the dataset description covers essential aspects such as the scope of data, query workloads, summary statistics, and possible tasks or applications.
+    A high score indicates that the description provides a comprehensive overview, including details on dataset size, structure, fields, and potential use cases.
+    2. Conciseness (1-10) - Measures the efficiency of the dataset description in conveying necessary information without redundancy.
+    A high score indicates that the description is succinct, avoiding unnecessary details while employing semantic types (e.g., categories, entities) to streamline communication.
+    3. Readability (1-10) -  Evaluates the logical flow and readability of the dataset description.
+    A high score suggests that the description progresses logically from one section to the next, creating a coherent and integrated narrative that facilitates understanding of the dataset.
+
+    Evaluation Steps:
+    Read the dataset description carefully and identify the main topic and key points. Assign a score for each criteria on a scale of 1 to 10, where 1 is the lowest and 10 is the highest based on the Evaluation Criteria.
+
+    Example 1:
+    Description: The dataset provides information on alcohol-impaired driving deaths and occupant deaths across various states in the United States. It includes data for 51 states, detailing the number of alcohol-impaired driving deaths and occupant deaths, with values ranging from 0 to 3723 and 0 to 10406, respectively. Each entry also contains the state abbreviation and its geographical coordinates. The dataset is structured with categorical and numerical data types, focusing on traffic safety and casualty statistics. Key attributes include state names, death counts, and location coordinates, making it a valuable resource for analyzing traffic safety trends and issues related to impaired driving.
+    Evaluation Form (scores ONLY): Completeness: 7, Conciseness: 9, Readability: 9
+
+    Example 2:
+    Description: The dataset provides a comprehensive overview of traffic safety statistics across various states in the United States, specifically focusing on alcohol-impaired driving deaths and occupant deaths. It includes data from 51 unique states, represented by their two-letter postal abbreviations, such as MA (Massachusetts), SD (South Dakota), AK (Alaska), MS (Mississippi), and ME (Maine). Each entry in the dataset captures critical information regarding the number of alcohol-impaired driving deaths and the total occupant deaths resulting from traffic incidents.
+    The column "Alcohol-Impaired Driving Deaths" is represented as an integer, indicating the number of fatalities attributed to alcohol impairment while driving. The dataset reveals a range of values, with the highest recorded number being 2367 deaths in Mississippi, highlighting the severity of the issue in certain regions. In contrast, states like Alaska report significantly lower figures, with only 205 alcohol-impaired driving deaths.
+    The "Occupant Deaths" column also consists of integer values, representing the total number of deaths among vehicle occupants, regardless of the cause. This data spans from 0 to 10406, with Mississippi again showing the highest number of occupant deaths at 6100, which raises concerns about overall traffic safety in the state.
+    Additionally, the dataset includes a "Location" column that provides geographical coordinates for each state, enhancing the spatial understanding of the data. The coordinates are formatted as latitude and longitude pairs, allowing for potential mapping and geographical analysis of traffic safety trends.
+    Overall, this dataset serves as a valuable resource for researchers, policymakers, and public safety advocates aiming to understand and address the impact of alcohol on driving safety across different states. It highlights the need for targeted interventions and policies to reduce alcohol-impaired driving incidents and improve occupant safety on the roads.
+    Evaluation Form (scores ONLY): Completeness: 8, Conciseness: 7, Readability: 8
+
+    Please provide scores for the given dataset description based on the Evaluation Criteria. Do not include any additional information or comments in your response.
+    Evaluation Form (scores ONLY):
+
+pairwise_evaluation:
+  system_message: |
+    You are a helpful and precise assistant for judging the quality of dataset descriptions.
+  user_prompt: |
+    You will be shown two descriptions of the same dataset, A and B.
+    Please select which description is better overall, based on completeness, conciseness, and readability.
+    Respond with exactly "A" or "B" (without any additional text).
+
+    Description A:
+    {desc_a}
+
+    Description B:
+    {desc_b}
+
+    Evaluation Criteria:
+    - Completeness (1-10) - Evaluates how thoroughly the dataset description covers essential aspects such as the scope of data, query workloads, summary statistics, and possible tasks or applications. A high score indicates that the description provides a comprehensive overview, including details on dataset size, structure, fields, and potential use cases.
+    - Conciseness (1-10) - Measures the efficiency of the dataset description in conveying necessary information without redundancy. A high score indicates that the description is succinct, avoiding unnecessary details while employing semantic types (e.g., categories, entities) to streamline communication.
+    - Readability (1-10) -  Evaluates the logical flow and readability of the dataset description. A high score suggests that the description progresses logically from one section to the next, creating a coherent and integrated narrative that facilitates understanding of the dataset.
+
+    Which description is better?

autoddg/description/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from .generator import DatasetDescriptionGenerator
+from .search import SearchFocusedDescription
+
+__all__ = ["DatasetDescriptionGenerator", "SearchFocusedDescription"]