Generate Dataframe Summaries With Python

langchain>=0.3.27
langchain-ollama>=0.3.7
pandas>=2.3.2

import pandas as pd
from langchain_ollama import ChatOllama
from typing import Literal

df = pd.read_csv("data/test.csv")

print("-*-" * 20)
print(f"Dataset shape: {df.shape}")
print("-*-" * 20)
print("Missing value stats:")
print(df.isnull().sum())

ollama run mistral

def get_llm(model_name: str = "mistral:latest") -> ChatOllama:
    """
    Create and configure a ChatOllama instance for local LLM inference.

    This function initializes a ChatOllama client configured to connect to a
    local Ollama server. The client is set up with deterministic output
    (temperature=0) for consistent responses across multiple calls with the
    same input.

    Parameters
    ----------
    model_name : str, optional
        The name of the Ollama model to use for chat completions.
        Must be a valid model name that is available on the local Ollama
        installation. Default is "mistral:latest".

    Returns
    -------
    ChatOllama
        A configured ChatOllama instance ready for chat completions.
    """
    return ChatOllama(
        model=model_name, base_url="http://localhost:11434", temperature=0
    )

print(get_llm().invoke("test").content)

def get_summary_context_message(df: pd.DataFrame, dataset_name:str) -> str:
    # Basic application statistics
    total_analisys = len(df)

    # Gender distribution
    gender_counts = df["Sex"].value_counts()
    male_count = gender_counts.get("M", 0)
    female_count = gender_counts.get("F", 0)

    # Stage Statistics
    stage_data = df["Stage"].dropna()
    stage_avg = stage_data.mean()
    stage_25th = stage_data.quantile(0.25)
    stage_50th = stage_data.quantile(0.50)
    stage_75th = stage_data.quantile(0.75)

    # NDays Statistics
    days_data = df["N_Days"].dropna()
    days_avg = days_data.mean()
    days_25th = days_data.quantile(0.25)
    days_50th = days_data.quantile(0.50)
    days_75th = days_data.quantile(0.75)

    def status_category(exp):
        if pd.isna(exp):
            return "Unkown"
        elif exp == "C":
            return "Censored"
        elif exp == "CL":
            return "Censored due to Lever tx"
        elif exp == "D":
            return "Death"
        else:
            return "Unkow"

    df['Status Str']= df['Status'].apply(status_category)
    status_str_stats = []

    for category in ["Censored", "Censored due to Lever tx", "Death",]:
        category_data = df[df["Status Str"] == category]
        if len(category_data) > 0:
            male = len(category_data[category_data["Sex"] == "M"])
            female = len(category_data[category_data["Sex"] == "F"])
            total = len(category_data)
            rate_m = (male / total) * 100
            rate_f = (female / total) * 100
            status_str_stats.append((category, male, female, total, rate_m, rate_f))

    summary =f"""{dataset_name}

Total Analisys: {total_analisys:,}

Gender Distribution:
- Male applicants: {male_count:,} ({male_count/total_analisys*100:.1f}%)
- Female applicants: {female_count:,} ({female_count/total_analisys*100:.1f}%)

Stage Statistics:
- Average Stage: {stage_avg:.2f}
- 25th percentile: {stage_25th:.2f}
- 50th percentile (median): {stage_50th:.2f}
- 75th percentile: {stage_75th:.2f}

N Day Statistics:
- N Days Stage: {days_avg:.2f}
- 25th percentile: {days_25th:.2f}
- 50th percentile (median): {days_50th:.2f}
- 75th percentile: {days_75th:.2f}
"""

    summary += "\n\nStatus Rates by Sex:"
    for category, male, female, total, rate_m, rate_f in status_str_stats:
        summary += (
            f"\n- {category}: {male}/{total} Male ({rate_m:.1f}% rate)"+
            f"\n- {category}: {female}/{total} Female ({rate_f:.1f}% rate)"

        )
    return summary

SUMMARIZE_DATAFRAME_PROMPT = """
You are an expert data analyst and data summarizer.
Your task is to take in complex datasets and return user-friendly descriptions and findings.

You were given this dataset:
- Name: {dataset_name}
- Source: {dataset_source}

This dataset was analyzed in a pipeline before it was given to you.
These are the findings returned by the analysis pipeline:

<context>
{context}
</context>

Based on these findings, write a detailed report in {report_format} format.
Give the report a meaningful title and separate findings into sections with headings and subheadings.
Output only the report in {report_format} and nothing else.

Report:
"""

def get_report_summary(
    dataset: pd.DataFrame,
    dataset_name: str,
    dataset_source: str,
    report_format: Literal["markdown", "html"] = "markdown",
) -> str:
    context_message = get_summary_context_message(df=dataset, dataset_name=dataset_name)
    prompt = SUMMARIZE_DATAFRAME_PROMPT.format(
        dataset_name=dataset_name,
        dataset_source=dataset_source,
        context=context_message,
        report_format=report_format,
    )
    return get_llm().invoke(input=prompt).content

md_report = get_report_summary(
    dataset=df,
    dataset_name="Cirrhosis Patient Survival Prediction",
    dataset_source="https://www.kaggle.com/datasets/joebeachcapital/cirrhosis-patient-survival-prediction/data"
)
print(md_report)