classifier_com.py

# Import built-in Python modules and third-party libraries

import os                     # Provides functions for interacting with the operating system
import re                     # Provides support for regular expressions (pattern matching)
import argparse               # Helps create command-line interfaces and parse arguments
import ollama                 # Ollama Python API for interacting with local LLMs
from tabulate import tabulate # Used to print formatted tables in terminal output


# -----------------------------------------------------------------------------
# FEW-SHOT EXAMPLES
# -----------------------------------------------------------------------------
# These examples are used in the Few-Shot prompting approach.
# The model is shown several examples of headlines and their correct labels
# before classifying the actual input headline.

FEW_SHOT_EXAMPLES = [

    # Politics examples
    {"text": "Senate Approves New Infrastructure Spending Bill.", "label": "Politics"},
    {"text": "Prime Minister Announces Snap Election Amid Coalition Collapse.", "label": "Politics"},
    {"text": "Governor Signs Executive Order Restricting Industrial Emissions.", "label": "Politics"},

    # Technology examples
    {"text": "New Quantum Computing Chip Breakthrough Multiplies Processing Speeds.", "label": "Technology"},
    {"text": "Cybersecurity Breach Exposes Millions of User Accounts Worldwide.", "label": "Technology"},
    {"text": "Startup Launches Revolutionary Virtual Reality Headset for Remote Workers.", "label": "Technology"},

    # Sports examples
    {"text": "Underdog Team Secures Dramatic Victory in Championship Final.", "label": "Sports"},
    {"text": "Olympic Committee Announces New Host City for Upcoming Summer Games.", "label": "Sports"},
    {"text": "World Number One Tennis Star Withdraws From Tournament Due to Injury.", "label": "Sports"},

    # Finance examples
    {"text": "Stock Market Plummets as Tech Sector Experiences Massive Sell-Off.", "label": "Finance"},
    {"text": "Global Conglomerate Reports Record Profits in Q3 Financial Release.", "label": "Finance"},
    {"text": "Cryptocurrency Regulations Tighten Across European Markets.", "label": "Finance"},

    # Health examples
    {"text": "Study Finds Regular Exercise Significantly Reduces Risk of Heart Disease.", "label": "Health"},
    {"text": "Hospitals Face Severe Nurse Shortages Amid Seasonal Flu Surge.", "label": "Health"},
    {"text": "Researchers Map Human Genome Sequence to Uncover Rare Genetic Mutations.", "label": "Health"}
]


# -----------------------------------------------------------------------------
# VALID CATEGORIES
# -----------------------------------------------------------------------------
# These are the only valid labels/categories allowed in classification.

VALID_CATEGORIES = [
    "Politics",
    "Technology",
    "Sports",
    "Finance",
    "Health"
]


# -----------------------------------------------------------------------------
# SYNTHETIC FALLBACK DATASET
# -----------------------------------------------------------------------------

SYNTHETIC_DATASET = [

    {
        "headline": "Tech Giants Agree on New Open-Source AI Safety Standards",
        "label": "Technology"
    },

    {
        "headline": "Central Bank Raises Interest Rates by 25 Basis Points to Combat Inflation",
        "label": "Finance"
    },

    {
        "headline": "Star Striker Signs Record-Breaking Five-Year Contract Extension",
        "label": "Sports"
    },

    {
        "headline": "Parliament Votes to Pass Historic Climate Action Bill After Fierce Debate",
        "label": "Politics"
    },

    {
        "headline": "New FDA-Approved Breakthrough Drug Shows Promise in Halting Alzheimer's",
        "label": "Health"
    },
]


# -----------------------------------------------------------------------------
# FUNCTION: clean_label
# -----------------------------------------------------------------------------
# Purpose:
# Extracts and validates the predicted category from raw LLM output.
#
# Why needed:
# Models may generate extra words/sentences instead of only the label.
# This function safely extracts the first valid category found.
#
# Parameters:
# raw_output (str) -> raw text generated by the model
#
# Returns:
# A valid category string OR "Unknown" if no valid category is found.

def clean_label(raw_output: str) -> str:

    # Loop through all valid categories
    for category in VALID_CATEGORIES:

        # Use regex search to check if the category exists in output
        # \b ensures whole-word matching
        # re.IGNORECASE ignores capitalization differences
        if re.search(rf"\b{category}\b", raw_output, re.IGNORECASE):

            # Return the matched valid category
            return category

    # If no valid category found, return "Unknown"
    return "Unknown"


# -----------------------------------------------------------------------------
# FUNCTION: classify_zero_shot
# -----------------------------------------------------------------------------
# Purpose:
# Performs Zero-Shot classification.
#
# Zero-Shot Prompting:
# The model is given ONLY instructions and no examples.
#
# Parameters:
# headline (str)   -> News headline to classify
# model_name (str) -> Ollama model name
#
# Returns:
# tuple:
#   (
#       predicted_label,
#       prompt_token_count,
#       generated_token_count
#   )

def classify_zero_shot(headline: str, model_name: str) -> tuple[str, int, int]:

    # Create system prompt that instructs the model
    system_prompt = (

        # Tell the model its role
        "You are an expert news editor. "

        # Explain the task
        "Your task is to classify the provided headline "

        # Show allowed categories
        f"into exactly one of these categories: {VALID_CATEGORIES}. "

        # Strict output formatting instructions
        "Respond with ONLY the category name. "
        "Do not write markdown, intros, or punctuation."
    )

    # Send chat request to Ollama
    response = ollama.chat(

        # Specify which model to use
        model=model_name,

        # Conversation messages
        messages=[

            # System role defines behavior/instructions
            {"role": "system", "content": system_prompt},

            # User message contains the actual headline
            {"role": "user", "content": f"Headline: {headline}"}
        ],

        # Deterministic output
        options={"temperature": 0.0}
    )

    # Extract generated text from response
    raw_content = response['message']['content'].strip()

    # Return:
    # cleaned label,
    # prompt token count,
    # generated token count
    return (
        clean_label(raw_content),
        response.get('prompt_eval_count', 0),
        response.get('eval_count', 0)
    )


# -----------------------------------------------------------------------------
# FUNCTION: classify_few_shot
# -----------------------------------------------------------------------------
# Purpose:
# Performs Few-Shot classification.
#
# Few-Shot Prompting:
# The model is shown multiple labeled examples before the real query.
#
# Parameters:
# headline (str)       -> headline to classify
# examples (list[dict])-> demonstration examples
# model_name (str)     -> Ollama model name
#
# Returns:
# tuple(label, prompt_tokens, generation_tokens)

def classify_few_shot(
    headline: str,
    examples: list[dict],
    model_name: str
) -> tuple[str, int, int]:

    # System prompt for few-shot prompting
    system_prompt = (

        # Define classification task
        f"Classify the input headline into one of these categories: {VALID_CATEGORIES}. "

        # Tell model to mimic example format
        "Follow the exact format shown in the examples. "

        # Restrict output format
        "Provide ONLY the category name."
    )

    # Convert each example into formatted text block
    example_blocks = [

        # Example format:
        # Headline: ...
        # Category: ...
        f"Headline: {ex['text']}\nCategory: {ex['label']}"

        # Iterate through all examples
        for ex in examples
    ]

    # Combine all examples into one user prompt
    # Then append actual query headline
    user_content = (

        # Separate examples using blank lines
        "\n\n".join(example_blocks)

        # Add actual headline to classify
        + f"\n\nHeadline: {headline}\nCategory:"
    )

    # Send request to Ollama
    response = ollama.chat(

        model=model_name,

        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content}
        ],

        # Deterministic generation
        options={"temperature": 0.0}
    )

    # Extract model response
    raw_content = response['message']['content'].strip()

    # Return cleaned prediction + token counts
    return (
        clean_label(raw_content),
        response.get('prompt_eval_count', 0),
        response.get('eval_count', 0)
    )


# -----------------------------------------------------------------------------
# FUNCTION: classify_cot
# -----------------------------------------------------------------------------
# Purpose:
# Performs Chain-of-Thought prompting.
#
# Chain-of-Thought (CoT):
# The model is encouraged to reason step-by-step before answering.
#
# Parameters:
# headline (str)
# model_name (str)
#
# Returns:
# tuple(label, prompt_tokens, generation_tokens)

def classify_cot(headline: str, model_name: str) -> tuple[str, int, int]:

    # CoT system prompt
    system_prompt = (

        # Task description
        f"Classify the news headline into one of: {VALID_CATEGORIES}. "

        # Ask the model to reason step-by-step
        "First, reason step-by-step about what fields or industries the words in the headline relate to. "

        # Force final answer format
        "Finally, end your response with the phrase "
        "'Final Category: <label>'."
    )

    # Query Ollama
    response = ollama.chat(

        model=model_name,

        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Headline: {headline}"}
        ],

        # Slightly higher temperature to encourage reasoning
        options={"temperature": 0.2}
    )

    # Extract raw generated output
    raw_content = response['message']['content'].strip()

    # Use regex to extract:
    # Final Category: <label>
    final_line_match = re.search(
        r"Final Category:\s*(\w+)",
        raw_content,
        re.IGNORECASE
    )

    # If explicit final category exists:
    if final_line_match:

        # Extract label from regex group
        label = clean_label(final_line_match.group(1))

    else:
        # Otherwise try cleaning whole response
        label = clean_label(raw_content)

    # Return prediction + token statistics
    return (
        label,
        response.get('prompt_eval_count', 0),
        response.get('eval_count', 0)
    )

# -----------------------------------------------------------------------------
# FUNCTION: load_data
# -----------------------------------------------------------------------------
# Purpose:
# Loads dataset from text files OR uses fallback synthetic dataset.
#
# Parameters:
# input_file  -> file containing headlines
# labels_file -> file containing labels
#
# Returns:
# list of dictionaries:
# [
#   {"headline": "...", "label": "..."},
#   ...
# ]

def load_data(input_file: str, labels_file: str) -> list[dict]:

    # If both files are provided
    if input_file and labels_file:

        try:

            # Open headline file
            with open(input_file, 'r', encoding='utf-8') as f_in, open(labels_file, 'r', encoding='utf-8') as f_lbl:

                    headlines = [
                        line.strip()
                        for line in f_in
                        if line.strip()
                    ]

                    # Read labels
                    labels = [
                        line.strip()
                        for line in f_lbl
                        if line.strip()
                    ]

            # Warn if counts mismatch
            if len(headlines) != len(labels):

                print(
                    f"Warning: Mismatch in lines! "
                    f"{len(headlines)} headlines vs "
                    f"{len(labels)} labels."
                )

            # Combine headlines and labels together
            return [
                {"headline": h, "label": l}
                for h, l in zip(headlines, labels)
            ]

        # Handle missing files
        except FileNotFoundError as e:

            print(f"Error loading files: {e}")

            # Exit program with error
            exit(1)

    else:

        # Fallback mode
        print(
            "No input files provided via arguments. "
            "Using built-in synthetic dataset."
        )

        return SYNTHETIC_DATASET


# -----------------------------------------------------------------------------
# FUNCTION: run_evaluation
# -----------------------------------------------------------------------------
# Purpose:
# Runs all prompting strategies on the dataset and computes metrics.
#
# Parameters:
# dataset (list[dict])
# model_name (str)

def run_evaluation(dataset: list[dict], model_name: str):

    # Startup information
    print(f"\nInitializing Prompt Engineering Playground...")
    print(
        f"Model: '{model_name}' | "
        f"Total Headlines: {len(dataset)}\n"
    )

    # Dictionary to track metrics for each prompting strategy
    metrics = {

        "Zero-Shot": {
            "correct": 0,
            "prompt_tok": 0,
            "comp_tok": 0
        },

        "Few-Shot": {
            "correct": 0,
            "prompt_tok": 0,
            "comp_tok": 0
        },

        "Chain-of-Thought": {
            "correct": 0,
            "prompt_tok": 0,
            "comp_tok": 0
        }
    }

    # Total dataset size
    total_items = len(dataset)

    # Iterate through dataset
    for idx, item in enumerate(dataset, 1):

        # Extract headline
        headline = item["headline"]

        # Extract true label
        ground_truth = item["label"]

        # Progress indicator
        print(f"Processing [{idx}/{total_items}]: \"{headline[:40]}...\"")

        # ---------------------------------------------------------------------
        # ZERO-SHOT
        # ---------------------------------------------------------------------

        # Run zero-shot classification
        zs_pred, zs_p_tok, zs_c_tok = classify_zero_shot(
            headline,
            model_name
        )

        # Accumulate token statistics
        metrics["Zero-Shot"]["prompt_tok"] += zs_p_tok
        metrics["Zero-Shot"]["comp_tok"] += zs_c_tok

        # Update accuracy count
        if zs_pred == ground_truth:
            metrics["Zero-Shot"]["correct"] += 1

        # ---------------------------------------------------------------------
        # FEW-SHOT
        # ---------------------------------------------------------------------

        fs_pred, fs_p_tok, fs_c_tok = classify_few_shot(
            headline,
            FEW_SHOT_EXAMPLES,
            model_name
        )

        metrics["Few-Shot"]["prompt_tok"] += fs_p_tok
        metrics["Few-Shot"]["comp_tok"] += fs_c_tok

        if fs_pred == ground_truth:
            metrics["Few-Shot"]["correct"] += 1

        # ---------------------------------------------------------------------
        # CHAIN-OF-THOUGHT
        # ---------------------------------------------------------------------

        cot_pred, cot_p_tok, cot_c_tok = classify_cot(
            headline,
            model_name
        )

        metrics["Chain-of-Thought"]["prompt_tok"] += cot_p_tok
        metrics["Chain-of-Thought"]["comp_tok"] += cot_c_tok

        if cot_pred == ground_truth:
            metrics["Chain-of-Thought"]["correct"] += 1

    # -------------------------------------------------------------------------
    # PRINT RESULTS
    # -------------------------------------------------------------------------

    print("\n" + "=" * 70)

    print("                 PROMPT DESIGN EVALUATION RESULTS                 ")

    print("=" * 70 + "\n")

    # Data structure for tabulate
    table_data = []

    # Process metrics for each approach
    for approach, data in metrics.items():

        # Accuracy percentage
        accuracy = (data["correct"] / total_items) * 100

        # Average prompt tokens
        avg_prompt = data["prompt_tok"] / total_items

        # Average generated tokens
        avg_comp = data["comp_tok"] / total_items

        # Total average tokens
        avg_total = avg_prompt + avg_comp

        # Add row to table
        table_data.append([
            approach,
            f"{data['correct']}/{total_items}",
            f"{accuracy:.1f}%",
            f"{avg_prompt:.1f}",
            f"{avg_comp:.1f}",
            f"{avg_total:.1f}"
        ])

    # Table headers
    headers = [
        "Approach",
        "Correct",
        "Accuracy",
        "Avg Prompt Tok",
        "Avg Gen Tok",
        "Avg Total Tok"
    ]

    # Print formatted table
    print(tabulate(
        table_data,
        headers=headers,
        tablefmt="grid"
    ))


# -----------------------------------------------------------------------------
# FUNCTION: main
# -----------------------------------------------------------------------------
# Entry point of the program.
#
# Responsibilities:
# - Parse CLI arguments
# - Validate arguments
# - Load dataset
# - Run evaluation

def main():

    # Create argument parser
    parser = argparse.ArgumentParser(

        # Program description
        description=(
            "News Classifier CLI: Compare Zero-Shot, "
            "Few-Shot, and CoT prompting."
        ),

        # Preserve formatting in help message
        formatter_class=argparse.RawTextHelpFormatter
    )

    # -------------------------------------------------------------------------
    # MODEL ARGUMENT
    # -------------------------------------------------------------------------

    parser.add_argument(

        # Short and long flag
        "-m", "--model",

        # Expected type
        type=str,

        # Default model
        default="llama3",

        # Help text
        help="Ollama model to use for classification (default: llama3)"
    )

    # -------------------------------------------------------------------------
    # INPUT FILE ARGUMENT
    # -------------------------------------------------------------------------

    parser.add_argument(

        "-i", "--input",

        type=str,

        help=(
            "Path to the plain text file containing "
            "news headlines (one per line)"
        )
    )

    # -------------------------------------------------------------------------
    # LABEL FILE ARGUMENT
    # -------------------------------------------------------------------------

    parser.add_argument(

        "-l", "--labels",

        type=str,

        help=(
            "Path to the plain text file containing "
            "ground truth labels (one per line)"
        )
    )

    # Parse command-line arguments
    args = parser.parse_args()

    # Validate:
    # Either BOTH input and labels are provided OR neither
    if bool(args.input) != bool(args.labels):

        parser.error(
            "You must provide both --input and --labels, "
            "or neither (to use fallback data)."
        )

    # Load dataset
    dataset = load_data(args.input, args.labels)

    # Run evaluation
    run_evaluation(dataset, args.model)


# -----------------------------------------------------------------------------
# PROGRAM ENTRY CHECK
# -----------------------------------------------------------------------------
# Ensures main() runs only if this file is executed directly.
# Prevents accidental execution when imported as a module.

if __name__ == "__main__":

    # Start program
    main()