Spaces:
Build error
Build error
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() |