mirror of
https://github.com/wassname/Castor.git
synced 2026-09-25 13:10:11 +08:00
rosequ
a471cea2c5
* reimplementation of SM model * features without normalization; parallel running of different modes * minor fix
39 lines
1.0 KiB
Python
39 lines
1.0 KiB
Python
from tqdm import tqdm
|
|
import array
|
|
import torch
|
|
import numpy as np
|
|
|
|
from argparse import ArgumentParser
|
|
|
|
|
|
def convert(fname, vocab):
|
|
save_file = '{}.pt'.format(fname)
|
|
stoi, vectors, dim = [], array.array('d'), None
|
|
|
|
# TODO: fix by reading the .dimensions file
|
|
vocab_size, dim = 2470719, 50
|
|
W = np.memmap(fname, dtype=np.double, shape=(vocab_size, dim))
|
|
|
|
|
|
print("Loading vectors from {}".format(fname))
|
|
vectors = []
|
|
for line in tqdm(W, total=len(W)):
|
|
entry = line
|
|
vectors.extend(entry)
|
|
|
|
vectors = torch.Tensor(vectors).view(-1, dim)
|
|
|
|
with open(vocab) as f:
|
|
stoi = {word.strip():i for i, word in enumerate(f)}
|
|
|
|
print('saving vectors to', save_file)
|
|
torch.save((stoi, vectors, dim), save_file)
|
|
|
|
if __name__ == '__main__':
|
|
parser = ArgumentParser(description='create word embedding')
|
|
parser.add_argument('--input', type=str, required=True)
|
|
parser.add_argument('--vocab', type=str, required=True)
|
|
|
|
args = parser.parse_args()
|
|
convert(args.input, args.vocab)
|