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() | |