print_hf_messages.py

from typing import Dict, List
from rich.console import Console
from rich.panel import Panel
from datasets import load_dataset

def print_hf_messages(messages: List[Dict[str, str]]):
    console = Console()
    colors = ["red", "green"]
    color_idx = 0
    console.rule(f"[bold yellow]The number of turns is {len(messages)}")
    for message in messages:
        role = message["role"]
        content = message["content"]
        console.print(Panel(content, title_align="left", title=role, border_style=colors[color_idx]))
        color_idx = (color_idx + 1) % 2

# Example usage
ds = load_dataset("allenai/tulu-v2-sft-mixture")
i = 0
while True:
    print_hf_messages(ds["train"][i]["messages"])
    i += 1
    input("Press Enter to continue...")