predict.py

import argparse
import csv
import json
import sys
from pathlib import Path

from deployment import load_bundle, load_networks, predict_text


def read_records(path, text_column):
    suffix = path.suffix.lower()

    if suffix == ".txt":
        records = []
        with path.open(encoding="utf-8") as input_file:
            for line_number, line in enumerate(input_file, start=1):
                text = line.strip()
                if text:
                    records.append({"id": line_number, text_column: text})
        return records

    if suffix == ".csv":
        with path.open(encoding="utf-8-sig", newline="") as input_file:
            records = list(csv.DictReader(input_file))
        if records and text_column not in records[0]:
            raise ValueError(
                f"CSV does not contain a '{text_column}' column. "
                f"Available columns: {list(records[0])}"
            )
        return records

    if suffix in {".jsonl", ".ndjson"}:
        records = []
        with path.open(encoding="utf-8") as input_file:
            for line_number, line in enumerate(input_file, start=1):
                if not line.strip():
                    continue
                record = json.loads(line)
                if not isinstance(record, dict):
                    raise ValueError(
                        f"JSONL line {line_number} must contain an object."
                    )
                if text_column not in record:
                    raise ValueError(
                        f"JSONL line {line_number} does not contain "
                        f"'{text_column}'."
                    )
                records.append(record)
        return records

    raise ValueError("Input file must be .txt, .csv, .jsonl, or .ndjson.")


def add_predictions(record, predictions):
    result = dict(record)
    for dimension, values in predictions.items():
        result[f"{dimension}_probability"] = values["probability"]
        result[f"{dimension}_prediction"] = values["prediction"]
    return result


def write_records(path, records):
    path.parent.mkdir(parents=True, exist_ok=True)
    suffix = path.suffix.lower()

    if suffix == ".csv":
        if not records:
            path.write_text("", encoding="utf-8")
            return
        with path.open("w", encoding="utf-8", newline="") as output_file:
            writer = csv.DictWriter(output_file, fieldnames=list(records[0]))
            writer.writeheader()
            writer.writerows(records)
        return

    if suffix in {".jsonl", ".ndjson"}:
        with path.open("w", encoding="utf-8") as output_file:
            for record in records:
                output_file.write(json.dumps(record) + "\n")
        return

    if suffix == ".json":
        path.write_text(json.dumps(records, indent=2), encoding="utf-8")
        return

    raise ValueError("Output file must be .csv, .json, .jsonl, or .ndjson.")


def main():
    parser = argparse.ArgumentParser(
        description="Predict personality dimensions from new text."
    )
    parser.add_argument("bundle", help="Path to a saved .pt model bundle.")
    parser.add_argument(
        "--text",
        help="A single text to classify.",
    )
    parser.add_argument(
        "--input",
        type=Path,
        help="Batch input file: .txt, .csv, .jsonl, or .ndjson.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Batch output file: .csv, .json, .jsonl, or .ndjson.",
    )
    parser.add_argument(
        "--text-column",
        default="text",
        help="Text field for CSV/JSONL input (default: text).",
    )
    args = parser.parse_args()

    bundle = load_bundle(args.bundle)
    networks = load_networks(bundle)

    if args.input:
        if args.text is not None:
            parser.error("--text and --input cannot be used together")
        if not args.output:
            parser.error("--output is required with --input")
        if not args.input.is_file():
            parser.error(f"input file does not exist: {args.input}")

        try:
            records = read_records(args.input, args.text_column)
            results = []
            for row_number, record in enumerate(records, start=1):
                text = record.get(args.text_column)
                if not isinstance(text, str) or not text.strip():
                    raise ValueError(
                        f"Row {row_number} has an empty or invalid "
                        f"'{args.text_column}' value."
                    )
                results.append(
                    add_predictions(
                        record,
                        predict_text(text, bundle, networks=networks),
                    )
                )
            write_records(args.output, results)
        except (OSError, ValueError, json.JSONDecodeError) as error:
            parser.error(str(error))

        print(f"Wrote {len(results)} prediction(s) to {args.output}")
        return

    if args.output:
        parser.error("--output can only be used with --input")

    text = args.text if args.text is not None else sys.stdin.read()
    if not text.strip():
        parser.error("prediction text cannot be empty")

    result = {
        "dataset": bundle["dataset"],
        "predictions": predict_text(text, bundle, networks=networks),
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()