| """Load and summarize IMDb reviews from local files.""" |
|
|
| from __future__ import annotations |
|
|
| import random |
| from collections import Counter |
| from pathlib import Path |
| from statistics import mean, median |
| from typing import Literal |
|
|
| from feedback_intelligence.types import ReviewRecord |
|
|
| Split = Literal["train", "test"] |
|
|
|
|
| def load_local_imdb_reviews( |
| base_path: Path, |
| split: Split, |
| sample_size: int | None = None, |
| seed: int = 42, |
| ) -> list[ReviewRecord]: |
| """Load deterministic review samples from a local IMDb directory.""" |
| categories = ("pos", "neg") |
| label_map = {"pos": "positive", "neg": "negative"} |
|
|
| records_by_label: dict[str, list[ReviewRecord]] = {} |
| rng = random.Random(seed) |
|
|
| for category in categories: |
| category_path = base_path / split / category |
| if not category_path.exists(): |
| raise FileNotFoundError(f"IMDb directory not found: {category_path}") |
|
|
| files = sorted(category_path.glob("*.txt")) |
| if not files: |
| raise FileNotFoundError(f"No review files found in: {category_path}") |
|
|
| if sample_size is None: |
| selected = files |
| else: |
| per_label = sample_size // 2 |
| if per_label <= 0: |
| raise ValueError("sample_size must be at least 2 for balanced sampling.") |
| if per_label > len(files): |
| raise ValueError( |
| f"Requested {per_label} files from {category_path}, " |
| f"but only {len(files)} are available." |
| ) |
| selected = rng.sample(files, per_label) |
|
|
| records_by_label[category] = [ |
| _record_from_path(path=file_path, split=split, label=label_map[category]) |
| for file_path in sorted(selected) |
| ] |
|
|
| records = records_by_label["pos"] + records_by_label["neg"] |
| rng.shuffle(records) |
| return records |
|
|
|
|
| def summarize_reviews(records: list[ReviewRecord]) -> dict[str, object]: |
| """Compute a compact dataset summary for CLI inspection.""" |
| if not records: |
| return { |
| "rows": 0, |
| "label_distribution": {}, |
| "word_count": {"min": 0, "median": 0, "mean": 0, "max": 0}, |
| } |
|
|
| word_counts = [record.word_count for record in records] |
| labels = Counter(record.label for record in records) |
| return { |
| "rows": len(records), |
| "label_distribution": dict(labels), |
| "word_count": { |
| "min": min(word_counts), |
| "median": int(median(word_counts)), |
| "mean": round(mean(word_counts), 2), |
| "max": max(word_counts), |
| }, |
| } |
|
|
|
|
| def _record_from_path(path: Path, split: str, label: str) -> ReviewRecord: |
| stem_parts = path.stem.split("_") |
| review_id = stem_parts[0] |
| score = int(stem_parts[1]) if len(stem_parts) == 2 else None |
| return ReviewRecord( |
| review_id=review_id, |
| text=path.read_text(encoding="utf-8"), |
| label=label, |
| split=split, |
| source="local-imdb", |
| score=score, |
| ) |
|
|
