from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import argparse
import os
import re
import sys
import fasttext
import fasttext.util
args = None
def eprint(*args, **kwargs):
print(*args, file=sys.stderr, **kwargs)
def guess_target_name(model_file, initial_dim, target_dim):
prg = re.compile("(.*).%s.(.*)" % initial_dim)
m = prg.match(model_file)
if m:
return "%s.%d.%s" % (m.group(1), target_dim, m.group(2))
sp_ext = os.path.splitext(model_file)
return "%s.%d%s" % (sp_ext[0], target_dim, sp_ext[1])
def command_reduce(model_file, target_dim, if_exists):
eprint("Loading model")
ft = fasttext.load_model(model_file)
initial_dim = ft.get_dimension()
if target_dim >= initial_dim:
raise Exception("Target dimension (%d) should be less than initial dimension (%d)." % (
target_dim, initial_dim))
result_filename = guess_target_name(model_file, initial_dim, target_dim)
if os.path.isfile(result_filename):
if if_exists == 'overwrite':
pass
elif if_exists == 'strict':
raise Exception(
"File already exists. Use --overwrite to overwrite.")
elif if_exists == 'ignore':
return result_filename
eprint("Reducing matrix dimensions")
fasttext.util.reduce_model(ft, target_dim)
eprint("Saving model")
ft.save_model(result_filename)
eprint("%s saved" % result_filename)
return result_filename
def main():
global args
parser = argparse.ArgumentParser(
description='fastText helper tool to reduce model dimensions.')
parser.add_argument("model", type=str,
help="model file to reduce. model.bin")
parser.add_argument("dim", type=int,
help="targeted dimension of word vectors.")
parser.add_argument("--overwrite", action="store_true",
help="overwrite if file exists.")
args = parser.parse_args()
command_reduce(args.model, args.dim, if_exists=(
'overwrite' if args.overwrite else 'strict'))
if __name__ == '__main__':
main()