File size: 1,010 Bytes
6fd648a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
import argparse
from termcolor import colored
from idiomifier import Idiomifier
from idiomify.fetchers import fetch_config, fetch_alpha
from transformers import BartTokenizer
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", type=str,
default="alpha")
parser.add_argument("--ver", type=str,
default="overfit")
parser.add_argument("--src", type=str,
default="If there's any benefits to losing my job, it's that I'll now be able to go to school full-time and finish my degree earlier.")
args = parser.parse_args()
config = fetch_config()[args.model][args.ver]
config.update(vars(args))
model = fetch_alpha(config['ver'])
tokenizer = BartTokenizer.from_pretrained(config['bart'])
idiomifier = Idiomifier(model, tokenizer)
src = config['src']
tgt = idiomifier(src=config['src'])
print(src, "\n->", colored(tgt, "blue"))
if __name__ == '__main__':
main()
|