idiomify / explore /explore_fetch_alpha_predict.py
eubinecto's picture
[#1] fetch_alpha implemented
322e083
raw
history blame
454 Bytes
from transformers import BartTokenizer
from builders import SourcesBuilder
from fetchers import fetch_alpha
def main():
model = fetch_alpha("overfit")
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()