File size: 462 Bytes
322e083
 
e3c7b5a
322e083
 
 
e3c7b5a
322e083
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from transformers import BartTokenizer
from builders import SourcesBuilder
from fetchers import fetch_idiomifier


def main():
    model = fetch_idiomifier("m-1-2")
    tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
    lit2idi = [
        ("my man", ""),
        ("hello", "")
    ]  # just some dummy stuff
    srcs = SourcesBuilder(tokenizer)(lit2idi)
    out = model.predict(srcs=srcs)
    print(out)


if __name__ == '__main__':
    main()