yangxinsci1993
edit interface
6e4fc64
raw
history blame contribute delete
No virus
867 Bytes
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()