from transformers import pipeline
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer
import argparse

# Load trained model
model = AutoModelForSeq2SeqLM.from_pretrained("output/reframer")
tokenizer = AutoTokenizer.from_pretrained("output/reframer")
reframer = pipeline('summarization', model=model, tokenizer=tokenizer)

def get_args():
    """ args from input
    """
    parser = argparse.ArgumentParser(description='HSIC-Bottleneck research')
    
    parser.add_argument('-ipt', '--input', required=True,
        type=str, help='input path')
    
    args = parser.parse_args()

    return args
    
def main():

    args = get_args()

    input_file = args.input

    with open(input_file, 'r') as file:
        data = file.read().rstrip()
        print(reframer(data)[0]['summary_text'])

if __name__ == '__main__':
    main()