citron-extract

#!/usr/bin/env python3
# Copyright 2021 BBC
# Authors: Chris Newell <chris.newell@bbc.co.uk>
#
# License: Apache-2.0

"""
This script runs Citron on the command line.

"""

import argparse
import logging
import json
import sys

from citron.citron import Citron
from citron.logger import logger


def main():
    parser = argparse.ArgumentParser(
        description="Extract quotes from text",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("-v",
        action = "store_true",
        default = False,
        help = "Verbose mode"
    )
    parser.add_argument("--model-path", 
        metavar = "model_path",
        type = str,
        required=True,
        help = "Path to Citron model directory"
    )
    parser.add_argument("--input-file",
        metavar = "input_file",
        type = str,
        help = "Optional: Otherwise read from stdin"
    )
    parser.add_argument("--output-file",
        metavar = "output_file",
        type = str,
        help = "Optional: Otherwise write to stdout"
    )
    args = parser.parse_args()
    
    if args.v:
        logger.setLevel(logging.DEBUG)
    
    citron = Citron(args.model_path)
    
    if args.input_file is None:
        text = ""
        
        while True:
            line = sys.stdin.readline()
            
            if not line:
                break
            
            text += " " + line
    
    else:
        with open(args.input_file, encoding="utf-8") as infile:
            text = infile.read()
    
    results = citron.extract(text)
    output = json.dumps(results, indent=4, sort_keys=False, ensure_ascii=False)
    
    if args.output_file is None:
        print(output)
    
    else:
        with open(args.output_file, "w", encoding="utf-8") as outfile:
            outfile.write(output + "\n")   


if __name__ == "__main__":
    main()