add charts galore

strickvl · strickvl · commit e1b10ce0c4b6 · 2024-10-31T11:12:04.000+01:00
diff --git a/llm-complete-guide/steps/populate_index.py b/llm-complete-guide/steps/populate_index.py
@@ -22,8 +22,9 @@
 import json
 import logging
 import math
-from typing import Annotated
+from typing import Annotated, Dict, List, Tuple
 
+from PIL import Image, ImageDraw, ImageFont
 from constants import (
     CHUNK_OVERLAP,
     CHUNK_SIZE,
@@ -40,12 +41,184 @@
 logger = logging.getLogger(__name__)
 
 
+def extract_docs_stats(total_documents: int, split_docs: List[Document]) -> Dict[str, Dict[str, int]]:
+    """Extracts statistics about the document chunks.
+
+    Args:
+        total_documents (int): The total number of original documents before splitting.
+        split_docs (List[Document]): The list of document chunks after splitting.
+
+    Returns:
+        Dict[str, Dict[str, int]]: A dictionary containing two sub-dictionaries:
+            - document_stats: Contains statistics about the chunks including:
+                - total_documents: Number of original documents
+                - total_chunks: Number of chunks after splitting
+                - avg_chunk_size: Average size of chunks in characters
+                - min_chunk_size: Size of smallest chunk in characters
+                - max_chunk_size: Size of largest chunk in characters
+            - chunks_per_section: Maps each document section to number of chunks it contains
+    """
+    total_documents = total_documents
+    total_chunks = len(split_docs)
+    chunk_sizes = [len(doc.page_content) for doc in split_docs]
+    avg_chunk_size = sum(chunk_sizes) / len(chunk_sizes)
+    min_chunk_size = min(chunk_sizes)
+    max_chunk_size = max(chunk_sizes)
+    chunks_per_section = {}
+    for doc in split_docs:
+        section = doc.parent_section
+        if section not in chunks_per_section:
+            chunks_per_section[section] = 0
+        chunks_per_section[section] += 1
+
+    return {
+        "document_stats": {
+            "total_documents": total_documents,
+            "total_chunks": total_chunks,
+            "avg_chunk_size": avg_chunk_size,
+            "min_chunk_size": min_chunk_size,
+            "max_chunk_size": max_chunk_size
+        },
+        "chunks_per_section": chunks_per_section
+    }
+
+
+
+def create_charts(stats: Dict[str, Dict[str, int]]) -> Image.Image:
+    """Creates a combined image containing a histogram of chunk sizes and a bar chart of chunk counts per section.
+
+    Args:
+        stats (Dict[str, Dict[str, int]]): A dictionary containing the extracted statistics.
+
+    Returns:
+        Image.Image: A combined image containing the histogram and bar chart.
+    """
+    document_stats = stats["document_stats"]
+    chunks_per_section = stats["chunks_per_section"]
+
+    # Create a new image with a white background
+    image_width = 800
+    image_height = 600
+    image = Image.new("RGB", (image_width, image_height), color="white")
+    draw = ImageDraw.Draw(image)
+
+    # Draw the histogram of chunk sizes
+    histogram_width = 400
+    histogram_height = 300
+    histogram_data = [document_stats["min_chunk_size"], document_stats["avg_chunk_size"], document_stats["max_chunk_size"]]
+    histogram_labels = ["Min", "Avg", "Max"]
+    histogram_x = 50
+    histogram_y = 50
+    draw_histogram(draw, histogram_x, histogram_y, histogram_width, histogram_height, histogram_data, histogram_labels)
+
+    # Draw the bar chart of chunk counts per section
+    bar_chart_width = 400
+    bar_chart_height = 300
+    bar_chart_data = list(chunks_per_section.values())
+    bar_chart_labels = list(chunks_per_section.keys())
+    bar_chart_x = 450
+    bar_chart_y = 50
+    draw_bar_chart(draw, bar_chart_x, bar_chart_y, bar_chart_width, bar_chart_height, bar_chart_data, bar_chart_labels)
+
+    # Add a title to the combined image
+    title_text = "Document Chunk Statistics"
+    title_font = ImageFont.truetype("arial.ttf", 24)
+    title_width, title_height = draw.textsize(title_text, font=title_font)
+    title_x = (image_width - title_width) // 2
+    title_y = 10
+    draw.text((title_x, title_y), title_text, font=title_font, fill="black")
+
+    return image
+
+def draw_histogram(draw: ImageDraw.Draw, x: int, y: int, width: int, height: int, data: List[int], labels: List[str]) -> None:
+    """Draws a histogram chart showing the distribution of chunk sizes.
+
+    Args:
+        draw (ImageDraw.Draw): The ImageDraw object to draw on
+        x (int): The x coordinate of the top-left corner of the histogram
+        y (int): The y coordinate of the top-left corner of the histogram 
+        width (int): The width of the histogram in pixels
+        height (int): The height of the histogram in pixels
+        data (List[int]): The values to plot in the histogram
+        labels (List[str]): The labels for each bar in the histogram
+    """
+    # Calculate the maximum value in the data
+    max_value = max(data)
+
+    # Calculate the bar width and spacing
+    bar_width = width // len(data)
+    bar_spacing = 10
+
+    # Draw the bars
+    for i, value in enumerate(data):
+        bar_height = (value / max_value) * height
+        bar_x = x + i * (bar_width + bar_spacing)
+        bar_y = y + height - bar_height
+        draw.rectangle([(bar_x, bar_y), (bar_x + bar_width, y + height)], fill="blue")
+
+        # Draw the label below the bar
+        label_text = labels[i]
+        label_font = ImageFont.truetype("arial.ttf", 12)
+        label_width, label_height = draw.textsize(label_text, font=label_font)
+        label_x = bar_x + (bar_width - label_width) // 2
+        label_y = y + height + 5
+        draw.text((label_x, label_y), label_text, font=label_font, fill="black")
+
+    # Draw the title above the histogram
+    title_text = "Chunk Size Distribution"
+    title_font = ImageFont.truetype("arial.ttf", 16)
+    title_width, title_height = draw.textsize(title_text, font=title_font)
+    title_x = x + (width - title_width) // 2
+    title_y = y - title_height - 10
+    draw.text((title_x, title_y), title_text, font=title_font, fill="black")
+
+def draw_bar_chart(draw: ImageDraw.Draw, x: int, y: int, width: int, height: int, data: List[int], labels: List[str]) -> None:
+    """Draws a bar chart showing the number of chunks per section.
+
+    Args:
+        draw (ImageDraw.Draw): The ImageDraw object to draw on
+        x (int): The x coordinate of the top-left corner of the bar chart
+        y (int): The y coordinate of the top-left corner of the bar chart
+        width (int): The width of the bar chart in pixels
+        height (int): The height of the bar chart in pixels
+        data (List[int]): The values to plot in the bar chart
+        labels (List[str]): The labels for each bar in the chart
+    """
+    # Calculate the maximum value in the data
+    max_value = max(data)
+
+    # Calculate the bar width and spacing
+    bar_width = width // len(data)
+    bar_spacing = 10
+
+    # Draw the bars
+    for i, value in enumerate(data):
+        bar_height = (value / max_value) * height
+        bar_x = x + i * (bar_width + bar_spacing)
+        bar_y = y + height - bar_height
+        draw.rectangle([(bar_x, bar_y), (bar_x + bar_width, y + height)], fill="green")
+
+        # Draw the label below the bar
+        label_text = labels[i]
+        label_font = ImageFont.truetype("arial.ttf", 12)
+        label_width, label_height = draw.textsize(label_text, font=label_font)
+        label_x = bar_x + (bar_width - label_width) // 2
+        label_y = y + height + 5
+        draw.text((label_x, label_y), label_text, font=label_font, fill="black")
+
+    # Draw the title above the bar chart
+    title_text = "Chunk Counts per Section"
+    title_font = ImageFont.truetype("arial.ttf", 16)
+    title_width, title_height = draw.textsize(title_text, font=title_font)
+    title_x = x + (width - title_width) // 2
+    title_y = y - title_height - 10
+    draw.text((title_x, title_y), title_text, font=title_font, fill="black")
+
 @step
 def preprocess_documents(
     documents: str,
-) -> Annotated[str, ArtifactConfig(name="split_chunks")]:
-    """
-    Preprocesses a JSON string of documents by splitting them into chunks.
+) -> Tuple[Annotated[str, ArtifactConfig(name="split_chunks")], Annotated[Image.Image, ArtifactConfig(name="doc_stats_chart")]]:
+    """Preprocesses a JSON string of documents by splitting them into chunks.
 
     Args:
         documents (str): A JSON string containing a list of documents to be preprocessed.
@@ -65,17 +238,22 @@ def preprocess_documents(
             },
         )
 
-        # Parse the JSON string into a list of Document objects
-        document_list = [Document(**doc) for doc in json.loads(documents)]
-
-        split_docs = split_documents(
+        document_list: List[Document] = [Document(**doc) for doc in json.loads(documents)]
+        split_docs: List[Document] = split_documents(
             document_list, chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP
         )
 
-        # Convert the list of Document objects back to a JSON string
-        split_docs_json = json.dumps([doc.__dict__ for doc in split_docs])
+        stats: Dict[str, Dict[str, int]] = extract_docs_stats(len(document_list), split_docs)
+        chart: Image.Image = create_charts(stats)
+
+        log_artifact_metadata(
+            artifact_name="split_chunks",
+            metadata=stats,
+        )
+
+        split_docs_json: str = json.dumps([doc.__dict__ for doc in split_docs])
 
-        return split_docs_json
+        return split_docs_json, chart
     except Exception as e:
         logger.error(f"Error in preprocess_documents: {e}")
         raise
