assignment.py - master - Timur11/vibe_coding_101-Timur11

assignment.py
312 строк10 KB
Timur11
update assignment.py
28 ноя 2025, 20:10
28 ноя 2025, 20:101d68b2b
'''
Analysis of SQuAD dataset (Stanford Question Answering Dataset).
Module provides functionality for comprehensive analysis of SQuAD dataset:
- Loading and processing question-answer dataset
- Statistical analysis of question, answer and context lengths
- Question classification by types (What, Who, When, etc.)
- Analysis of context text uniqueness and distribution
- Visualization of analysis results
- Saving results in JSON format

Usage:
    python squad_analysis.py

Output files:
    - squad_results.json: complete analysis results
    - squad_distribution.png: length distribution visualizations
'''

import json
from collections import Counter
from typing import Any, Dict

import matplotlib.pyplot as plt
import numpy as np
from datasets import DatasetDict, load_dataset

# Constants
DPI = 300
OUTPUT_IMAGE = 'squad_distribution.png'
OUTPUT_JSON = 'squad_results.json'


def load_squad_dataset() -> DatasetDict:
    """Load SQuAD dataset from Hugging Face.

    Returns:
        DatasetDict: Loaded SQuAD dataset with train and validation splits
    """
    return load_dataset('squad')


def analyze_qa_statistics(dataset: DatasetDict) -> Dict[str, Any]:
    """Analyze statistics of question, answer and context lengths.

    Args:
        dataset: Loaded SQuAD dataset

    Returns:
        Dict with statistics of question, answer and context lengths
    """
    train_data = dataset['train']

    question_lengths = [
        len(question.split())
        for question in train_data['question']
    ]

    answer_lengths = [
        len(answer['text'][0].split())
        for answer in train_data['answers']
    ]

    context_lengths = [
        len(context.split())
        for context in train_data['context']
    ]

    return {
        'questions': {
            'mean_length': float(np.mean(question_lengths)),
            'median_length': float(np.median(question_lengths)),
            'std_length': float(np.std(question_lengths)),
            'min_length': int(np.min(question_lengths)),
            'max_length': int(np.max(question_lengths))
        },
        'answers': {
            'mean_length': float(np.mean(answer_lengths)),
            'median_length': float(np.median(answer_lengths)),
            'std_length': float(np.std(answer_lengths)),
            'min_length': int(np.min(answer_lengths)),
            'max_length': int(np.max(answer_lengths))
        },
        'contexts': {
            'mean_length': float(np.mean(context_lengths)),
            'median_length': float(np.median(context_lengths)),
            'std_length': float(np.std(context_lengths)),
            'min_length': int(np.min(context_lengths)),
            'max_length': int(np.max(context_lengths))
        }
    }


def classify_questions(dataset: DatasetDict) -> Dict[str, int]:
    """Classify questions by types based on first words.

    Args:
        dataset: Loaded SQuAD dataset

    Returns:
        Dict with question type distribution by categories
    """
    train_questions = dataset['train']['question']
    question_types = Counter()

    for question in train_questions:
        if not question or not question.strip():
            question_types['Other'] += 1
            continue

        words = question.strip().split()
        if not words:
            question_types['Other'] += 1
            continue

        first_word = words[0].lower()

        if first_word.startswith('what'):
            question_types['What'] += 1
        elif first_word.startswith('who'):
            question_types['Who'] += 1
        elif first_word.startswith('when'):
            question_types['When'] += 1
        elif first_word.startswith('where'):
            question_types['Where'] += 1
        elif first_word.startswith('why'):
            question_types['Why'] += 1
        elif first_word.startswith('how'):
            question_types['How'] += 1
        elif first_word.startswith('which'):
            question_types['Which'] += 1
        elif first_word.startswith('whose'):
            question_types['Whose'] += 1
        elif first_word.startswith('whom'):
            question_types['Whom'] += 1
        else:
            question_types['Other'] += 1

    return dict(question_types)


def analyze_context_lengths(dataset: DatasetDict) -> Dict[str, Any]:
    """Analyze context texts - extended statistics.

    Args:
        dataset: Loaded SQuAD dataset

    Returns:
        Dict with context statistics
    """
    train_contexts = dataset['train']['context']

    context_lengths = [len(context.split()) for context in train_contexts]
    char_lengths = [len(context) for context in train_contexts]

    unique_contexts = len(set(train_contexts))
    total_contexts = len(train_contexts)

    duplicate_ratio = round(
        (total_contexts - unique_contexts) / total_contexts * 100, 2
    )

    return {
        'word_statistics': {
            'mean': float(np.mean(context_lengths)),
            'median': float(np.median(context_lengths)),
            'std': float(np.std(context_lengths)),
            'min': int(np.min(context_lengths)),
            'max': int(np.max(context_lengths))
        },
        'char_statistics': {
            'mean': float(np.mean(char_lengths)),
            'median': float(np.median(char_lengths)),
            'std': float(np.std(char_lengths)),
            'min': int(np.min(char_lengths)),
            'max': int(np.max(char_lengths))
        },
        'uniqueness': {
            'unique_contexts': unique_contexts,
            'total_contexts': total_contexts,
            'duplicate_ratio': duplicate_ratio
        }
    }


def create_distribution_chart(dataset: DatasetDict, stats: Dict[str, Any]) -> None:
    """Create distribution charts for lengths.

    Args:
        dataset: Loaded SQuAD dataset
        stats: Statistics from analyze_qa_statistics
    """
    train_data = dataset['train']

    question_lengths = [len(q.split()) for q in train_data['question']]
    answer_lengths = [len(a['text'][0].split()) for a in train_data['answers']]
    context_lengths = [len(c.split()) for c in train_data['context']]

    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(15, 12))

    # Question histogram
    ax1.hist(question_lengths, bins=50, alpha=0.7, color='skyblue', edgecolor='black')
    ax1.set_title('Distribution of Question Lengths (Words)')
    ax1.set_xlabel('Number of Words')
    ax1.set_ylabel('Frequency')
    mean_q = stats['questions']['mean_length']
    median_q = stats['questions']['median_length']
    ax1.axvline(mean_q, color='red', linestyle='--', label=f'Mean: {mean_q:.1f}')
    ax1.axvline(median_q, color='green', linestyle='--', label=f'Median: {median_q:.1f}')
    ax1.legend()

    # Answer histogram
    ax2.hist(answer_lengths, bins=50, alpha=0.7, color='lightcoral', edgecolor='black')
    ax2.set_title('Distribution of Answer Lengths (Words)')
    ax2.set_xlabel('Number of Words')
    ax2.set_ylabel('Frequency')
    mean_a = stats['answers']['mean_length']
    median_a = stats['answers']['median_length']
    ax2.axvline(mean_a, color='red', linestyle='--', label=f'Mean: {mean_a:.1f}')
    ax2.axvline(median_a, color='green', linestyle='--', label=f'Median: {median_a:.1f}')
    ax2.legend()

    # Context histogram
    ax3.hist(context_lengths, bins=50, alpha=0.7, color='lightgreen', edgecolor='black')
    ax3.set_title('Distribution of Context Lengths (Words)')
    ax3.set_xlabel('Number of Words')
    ax3.set_ylabel('Frequency')
    mean_c = stats['contexts']['mean_length']
    median_c = stats['contexts']['median_length']
    ax3.axvline(mean_c, color='red', linestyle='--', label=f'Mean: {mean_c:.1f}')
    ax3.axvline(median_c, color='green', linestyle='--', label=f'Median: {median_c:.1f}')
    ax3.legend()

    # Box plot for comparison
    all_lengths = [question_lengths, answer_lengths, context_lengths]
    ax4.boxplot(all_lengths, labels=['Questions', 'Answers', 'Contexts'])
    ax4.set_title('Comparison of Length Distributions')
    ax4.set_ylabel('Number of Words')
    ax4.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig(OUTPUT_IMAGE, dpi=DPI, bbox_inches='tight')
    plt.show()


def save_results_to_json(
    qa_stats: Dict[str, Any],
    question_types: Dict[str, int],
    context_stats: Dict[str, Any],
    dataset_info: Dict[str, Any],
) -> None:
    """Save all analysis results to JSON file.

    Args:
        qa_stats: Question-answer statistics
        question_types: Question type distribution
        context_stats: Context statistics
        dataset_info: Dataset information
    """
    most_common_type = max(question_types.items(), key=lambda x: x[1])[0]

    results = {
        'dataset_info': dataset_info,
        'qa_statistics': qa_stats,
        'question_type_distribution': question_types,
        'context_analysis': context_stats,
        'summary': {
            'total_questions': dataset_info['train_size'],
            'most_common_question_type': most_common_type,
            'avg_question_length': qa_stats['questions']['mean_length'],
            'avg_answer_length': qa_stats['answers']['mean_length'],
            'avg_context_length': qa_stats['contexts']['mean_length'],
        },
    }

    with open(OUTPUT_JSON, 'w', encoding='utf-8') as file:
        json.dump(results, file, indent=2, ensure_ascii=False)


def main() -> None:
    """Main function for complete SQuAD dataset analysis."""
    print('🔄 Loading SQuAD dataset...')
    dataset = load_squad_dataset()

    dataset_info = {
        'train_size': len(dataset['train']),
        'validation_size': len(dataset['validation']),
        'features': list(dataset['train'].features.keys()),
    }

    print('📊 Analyzing question-answer statistics...')
    qa_stats = analyze_qa_statistics(dataset)

    print('🏷️ Classifying questions by types...')
    question_types = classify_questions(dataset)

    print('📖 Analyzing context texts...')
    context_stats = analyze_context_lengths(dataset)

    print('📈 Creating visualizations...')
    create_distribution_chart(dataset, qa_stats)

    print('💾 Saving results...')
    save_results_to_json(qa_stats, question_types, context_stats, dataset_info)

    print('✅ Analysis completed successfully!')
    print('📁 Created files:')
    print(f'   - {OUTPUT_JSON}')
    print(f'   - {OUTPUT_IMAGE}')


if __name__ == '__main__':
    main()