File size: 859 Bytes
08409ff
 
 
6fd648a
08409ff
c1728bd
6fd648a
 
 
 
 
08409ff
c1728bd
6fd648a
642d911
6fd648a
e3c7b5a
c1728bd
e3c7b5a
642d911
08409ff
c1728bd
 
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
"""
This is for just a simple sanity check on the inference.
"""
import argparse
from idiomify.pipeline import Pipeline
from idiomify.fetchers import fetch_config, fetch_idiomifier, fetch_tokenizer
from transformers import BartTokenizer


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sent", type=str,
                        default="Just remember that there will always be a hope even when things look hopeless")
    args = parser.parse_args()
    config = fetch_config()['idiomifier']
    config.update(vars(args))
    model = fetch_idiomifier(config['ver'])
    tokenizer = fetch_tokenizer(config['tokenizer_ver'])
    model.eval()  # this is crucial
    pipeline = Pipeline(model, tokenizer)
    src = config['sent']
    tgts = pipeline(sents=[src])
    print(src, "\n->", tgts[0])


if __name__ == '__main__':
    main()