bench.py

import json
import logging
import sys
import time
from pathlib import Path

from gensim.models.fasttext import FastText
from gensim.models.word2vec import Text8Corpus


def train(target_dir):
    corpus = Text8Corpus('./text8')
    model = FastText(corpus)
    model.save(str(target_dir/'model'))


def evaluate(model_dir):
    model = FastText.load(str(model_dir/'model'))
    accuracy = model.accuracy('./questions-words.txt')
    with (model_dir/'accuracy.json').open('wt') as fp:
        json.dump(accuracy, fp)

if __name__ == '__main__':
    if len(sys.argv) != 3:
        print('bench.py [train|evaluate] target_dir', file=sys.stderr)
        sys.exit(1)
    target_dir = Path(sys.argv[2])
    if not target_dir.exists():
        target_dir.mkdir()
    logging.basicConfig(level=logging.INFO)
    time.sleep(15)
    if sys.argv[1] == 'train':
        train(Path(target_dir))
    elif sys.argv[1] == 'evaluate':
        evaluate(Path(target_dir))