Initial Commit

README.md

@@ -1,10 +1,10 @@
 ---
 title: Dataset Test Workflow
-emoji: 💬
+emoji: 🧪
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 4.42.0
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,63 +1,234 @@
+import os
+import json
+import duckdb
 import gradio as gr
+import pandas as pd
+import pandera as pa
+from pandera import Column
+import ydata_profiling as pp
 from huggingface_hub import InferenceClient
+from prompt import PROMPT_PANDERA
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
+# Height of the Tabs Text Area
+TAB_LINES = 8
+# Load Token
+md_token = os.getenv('MD_TOKEN')
+os.environ['HF_TOKEN'] = os.getenv('HF_TOKEN')
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+INPUT_PROMPT = '''
+Here is the frist few samples of data:
+<Sample Data>
+{data}
+</Sample Data<>
+'''
 
-    messages.append({"role": "user", "content": message})
 
-    response = ""
+print('Connecting to DB...')
+# Connect to DB
+conn = duckdb.connect(f"md:my_db?motherduck_token={md_token}", read_only=True)
+client = InferenceClient("meta-llama/Meta-Llama-3-70B-Instruct")
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+# Get Databases
+def get_schemas():
+    schemas = conn.execute("""
+    SELECT DISTINCT schema_name
+    FROM information_schema.schemata
+    WHERE schema_name NOT IN ('information_schema', 'pg_catalog')
+    """).fetchall()
+    return [item[0] for item in schemas]
 
-        response += token
-        yield response
+# Get Tables
+def get_tables_names(schema_name):
+    tables = conn.execute(f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{schema_name}'").fetchall()
+    return [table[0] for table in tables]
 
+# Update Tables
+def update_table_names(schema_name):
+    tables = get_tables_names(schema_name)
+    return gr.update(choices=tables)
+
+def get_data_df(schema):
+    print('Getting Dataframe from the Database')
+    return conn.sql(f"SELECT * FROM {schema} LIMIT 1000").df()
+
+def run_llm(df):
+  messages=[
+        {"role": "system", "content": PROMPT_PANDERA},
+        {"role": "user", "content": INPUT_PROMPT.format(data=df.head().to_json(orient='records'))},
+    ]
+  try:
+    response = client.chat_completion(messages, max_tokens=1024)
+    print(response.choices[0].message.content)
+    tests = json.loads(response.choices[0].message.content)
+  except Exception as e:
+      return e
+  return tests
+
+# Get Schema
+def get_table_schema(table):
+    result = conn.sql(f"SELECT sql, database_name, schema_name FROM duckdb_tables() where table_name ='{table}';").df()
+    ddl_create = result.iloc[0,0]
+    parent_database = result.iloc[0,1]
+    schema_name = result.iloc[0,2]
+    full_path = f"{parent_database}.{schema_name}.{table}"
+    if schema_name != "main":
+        old_path = f"{schema_name}.{table}"
+    else:
+        old_path = table
+    ddl_create = ddl_create.replace(old_path, full_path)
+    return full_path
+
+def describe(df):
+    numerical_info = df.select_dtypes(include=['number']).describe().T.reset_index()
+    numerical_info.rename(columns={'index': 'column'}, inplace=True)
+    
+    categorical_info = df.select_dtypes(include=['object']).describe().T.reset_index()
+    categorical_info.rename(columns={'index': 'column'}, inplace=True)
+
+    return numerical_info, categorical_info
+
+def validate_pandera(tests, df):
+    validation_results = []
+
+    # Loop through each test rule and validate each column separately
+    for test in tests:
+        column_name = test['column_name']
+        rule = eval(test['pandera_rule'])  # Evaluate the Pandera column rule
+
+        try:
+            # Apply the rule to the column and validate
+            validated_column = rule(df[[column_name]])  # Validate the specific column
+            validation_results.append({
+            "Columns": column_name,
+            "Result": "✅ Pass"
+            })
+        except Exception as e:
+            # If validation fails, catch the exception and mark the column as 'Fail'
+            validation_results.append({
+            "Columns": column_name,
+            "Result": f"❌ Fail - {str(e)}"
+            })
+    return pd.DataFrame(validation_results)
+
+def statistics(df):
+    profile = pp.ProfileReport(df)
+    report_dict = profile.get_description()
+    description, alerts = report_dict.table, report_dict.alerts
+    # Statistics
+    mapping = {
+        'n': 'Number of observations',
+        'n_var': 'Number of variables',
+        'n_cells_missing': 'Number of cells missing',
+        'n_vars_with_missing': 'Number of columns with missing data',
+        'n_vars_all_missing': 'Columns with all missing data',
+        'p_cells_missing': 'Missing cells (%)',
+        'n_duplicates': 'Duplicated rows',
+        'p_duplicates': 'Duplicated rows (%)',
+    }
+
+    updated_data = {mapping.get(k, k): v for k, v in description.items() if k != 'types'}
+    # Add flattened types information
+    if 'Text' in description.get('types', {}):
+            updated_data['Number of text columns'] = description['types']['Text']
+    if 'Categorical' in description.get('types', {}):
+        updated_data['Number of categorical columns'] = description['types']['Categorical']
+    if 'Numeric' in description.get('types', {}):
+        updated_data['Number of numeric columns'] = description['types']['Numeric']
+    if 'DateTime' in description.get('types', {}):
+        updated_data['Number of datetime columns'] = description['types']['DateTime']
+
+    df_statistics = pd.DataFrame(list(updated_data.items()), columns=['Statistic Description', 'Value'])
+    df_statistics['Value'] = df_statistics['Value'].astype(int)
+
+    # Alerts
+    alerts_list = [(str(alert).replace('[', '').replace(']', ''), alert.alert_type_name) for alert in alerts]
+    df_alerts = pd.DataFrame(alerts_list, columns=['Data Quality Issue', 'Category'])
+
+    return df_statistics, df_alerts
+# Main Function
+def main(table):
+    schema = get_table_schema(table)
+    df = get_data_df(schema)
+    df_statistics, df_alerts = statistics(df)
+    describe_cat, describe_num = describe(df)
+
+    tests = run_llm(df)
+    print(tests)
+    if isinstance(tests, Exception):
+        tests = pd.DataFrame([{"error": f"❌ Unable to get the SQL query based on the text. {tests}"}])
+        return df.head(10), df_statistics, df_alerts, describe_cat, describe_num, tests, pd.DataFrame([])
+
+    tests_df = pd.DataFrame(tests)
+    tests_df.rename(columns={tests_df.columns[0]: 'Column', tests_df.columns[1]: 'Rule Name', tests_df.columns[2]: 'Rules' }, inplace=True)
+    pandera_results = validate_pandera(tests, df)
+    
+    return df.head(10), df_statistics, df_alerts, describe_cat, describe_num, tests_df, pandera_results
+
+# Custom CSS styling
+custom_css = """
+.gradio-container {
+    background-color: #f0f4f8;
+}
+.logo {
+    max-width: 200px;
+    margin: 20px auto;
+    display: block;
+}
+.gr-button {
+    background-color: #4a90e2 !important;
+}
+.gr-button:hover {
+    background-color: #3a7bc8 !important;
+}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="indigo"), css=custom_css) as demo:
+    gr.Image("logo.png", label=None, show_label=False, container=False, height=100)
+
+    gr.Markdown("""
+    <div style='text-align: center;'>
+    <strong style='font-size: 36px;'>Dataset Test Workflow</strong>
+    <br>
+    <span style='font-size: 20px;'>Implement and Automate Data Validation Processes.</span>
+    </div>
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            schema_dropdown = gr.Dropdown(choices=get_schemas(), label="Select Schema", interactive=True)
+            tables_dropdown = gr.Dropdown(choices=[], label="Available Tables", value=None)
+            with gr.Row():
+                generate_query_button = gr.Button("Validate Data", variant="primary")
+
+        with gr.Column(scale=2):
+            with gr.Tabs():
+
+                with gr.Tab("Description"):
+                    with gr.Row():
+                        with gr.Column(min_width=220):
+                            data_description = gr.DataFrame(label="Data Description", value=[], interactive=False)
+                    with gr.Row():
+                        with gr.Column(min_width=320):
+                            describe_cat = gr.DataFrame(label="Categorical Information", value=[], interactive=False)
+                        with gr.Column(min_width=320):    
+                            describe_num = gr.DataFrame(label="Numerical Information", value=[], interactive=False)
+
+                with gr.Tab("Alerts"):
+                    data_alerts = gr.DataFrame(label="Alerts", value=[], interactive=False)
+
+                with gr.Tab("Rules & Validations"):
+                    tests_output = gr.DataFrame(label="Validation Rules", value=[], interactive=False)
+                    test_result_output = gr.DataFrame(label="Validation Result", value=[], interactive=False)
+                
+                with gr.Tab("Data"):
+                    result_output = gr.DataFrame(label="Dataframe (10 Rows)", value=[], interactive=False)
+
+        schema_dropdown.change(update_table_names, inputs=schema_dropdown, outputs=tables_dropdown)
+        generate_query_button.click(main, inputs=[tables_dropdown], outputs=[result_output, data_description, data_alerts, describe_cat, describe_num, tests_output, test_result_output])
+
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True)

prompt.py

@@ -0,0 +1,64 @@
+PROMPT_PANDERA = """
+You are a data quality engineer. Your role is to create deterministic rules to validate the quality of a dataset using **Pandera**.
+You will be provided with the first few rows of data below that represents the dataset for which you need to create validation rules. Please note that this is only a sample of the data, and there may be additional rows and categorical columns that are not fully represented in the sample. Keep in mind that the sample may not cover all possible values, but the validation rules must handle all data in the dataset.
+
+Follow this process:
+
+1. **Observe the sample data.**
+2. **For each column**, create a validation rule using Pandera syntax.
+    Here are the valid pandera check class methods DO NOT USE ANYOTHER METHODS OTHER THAN THE BELOW GIVEN METHODS:
+    [
+    'pa.Check.between(min_value, max_value, include_min=True, include_max=True, **kwargs)',
+    'pa.Check.eq(value, **kwargs)',
+    'pa.Check.equal_to(value, **kwargs)',
+    'pa.Check.ge(min_value, **kwargs)',
+    'pa.Check.greater_than(min_value, **kwargs)',
+    'pa.Check.greater_than_or_equal_to(min_value, **kwargs)',
+    'pa.Check.gt(min_value, **kwargs)',
+    'pa.Check.in_range(min_value, max_value, include_min=True, include_max=True, **kwargs)',
+    'pa.Check.isin(allowed_values, **kwargs)',
+    'pa.Check.le(max_value, **kwargs)',
+    'pa.Check.less_than(max_value, **kwargs)',
+    'pa.Check.less_than_or_equal_to(max_value, **kwargs)',
+    'pa.Check.lt(max_value, **kwargs)',
+    'pa.Check.ne(value, **kwargs)',
+    'pa.Check.not_equal_to(value, **kwargs)',
+    'pa.Check.notin(forbidden_values, **kwargs)',
+    'pa.Check.str_contains(pattern, **kwargs)',
+    'pa.Check.str_endswith(string, **kwargs)',
+    'pa.Check.str_length(min_value=None, max_value=None, **kwargs)',
+    'pa.Check.str_matches(pattern, **kwargs)',
+    'pa.Check.str_startswith(string, **kwargs)',
+    'pa.Check.unique_values_eq(values, **kwargs)'
+    ]
+    ALSO DONT USE REGEX FOR VALIDATIONS
+3. Ensure that each rule specifies the expected data type and applies necessary checks such as:
+      name argument should be a valid column name. DO NOT USE ANYOTHER PANDERA 
+   - **Data Type Validation** (e.g., `pa.Column(int, nullable=False, name="age")` ensures integers)
+   - **Non-null Check** (e.g., `pa.Column(str, nullable=False, name="name")` to ensure no nulls are allowed)
+   - **Unique Value Check** (e.g., `pa.Column(int, unique=True, name="ID")` for uniqueness)
+   - **Range or Bound Checks** (e.g., `pa.Column(float, checks=pa.Check.in_range(min_value=0, max_value=100), name="score")` for numerical ranges)
+   - **Allowed Value Checks** (e.g., `pa.Column(str, checks=pa.Check.isin([value1, value2]), name="gender")` to restrict values to a set)
+   - **Custom Validation Logic** using `pa.Column(int, checks=pa.Check(lambda x: x % 2 == 0), name="even_number")` with lambda functions (e.g., custom logic for even numbers or string patterns)
+  FOR DATETIME OR DATE COLUMN USE THE BELOW VALIDATION DO NOT CONISER IT AS INT OR FLOAT
+   - **DateTime or Date Validation** (e.g., `pa.Column(pa.dtypes.Timestamp, nullable=False), name="date_column")` to ensure dates or datetime)
+
+   For each column, provide a **column name**, **rule name** and a pandera_rule. Example structure:
+
+   ```json
+   [
+     {
+       "column_name": "age",
+       "rule_name": "Ensure Column is Integer",
+       "pandera_rule": "Column(int, nullable=False, name='age')"
+     },
+     {
+       "column_name": "ID",
+       "rule_name": "Unique Identifier Check",
+       "pandera_rule": "Column(int, unique=True, name='ID')"
+     }
+   ]
+3 Repeat this process for max 5 columns in the dataset. If the data is less than 5 columns than include all columns. Group all the rules into a single JSON object and ensure that there is at least one validation rule for each column.
+Return the final rules as a single JSON object, ensuring that each column is thoroughly validated based on the observations of the sample data.
+DO NOT RETURN ANYTHING OR ANY EXPLAINATION OTHER THAN JSON
+"""

requirements.txt

@@ -1 +1,7 @@
-huggingface_hub==0.22.2
+torch
+huggingface_hub
+accelerate
+transformers
+duckdb
+pandera
+ydata-profiling