idiomify / explore /explore_fetch_idiomifier_predict.py
eubinecto's picture
[#2] Support for d-1-2 done. Support for m-1-2 partially done (need to implement the testing logic with some valid metrics)
e3c7b5a
raw
history blame contribute delete
462 Bytes
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()