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