mirror of
https://github.com/wassname/Castor.git
synced 2026-09-26 13:20:14 +08:00
+ cleaned up the model code for the simple qa directory + created vocab objects for pre-loading word embeddings easily
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
import torch
|
|
import nltk
|
|
import string
|
|
from torch.autograd import Variable
|
|
import numpy as np
|
|
|
|
## functions for loading data from disk
|
|
|
|
def process_tokenize_text(text):
|
|
punc_remover = str.maketrans('', '', string.punctuation)
|
|
processed_text = text.lower().translate(punc_remover)
|
|
tokens = nltk.word_tokenize(processed_text)
|
|
return tokens
|
|
|
|
def read_embedding(embed_pt_filepath):
|
|
embed_tuple = torch.load(embed_pt_filepath)
|
|
word2index, w2v_tensor, dim = embed_tuple
|
|
return word2index, w2v_tensor
|
|
|
|
|
|
def find_max_seq_length(text):
|
|
max_len = -1
|
|
for tokens in text:
|
|
curr_len = len(tokens)
|
|
if curr_len > max_len:
|
|
max_len = curr_len
|
|
return max_len
|
|
|
|
|
|
def read_text_tensor(batch_text, word_vocab):
|
|
out_text = []
|
|
max_len = find_max_seq_length(batch_text)
|
|
for sent_tokens in batch_text:
|
|
S = len(sent_tokens)
|
|
sent = []
|
|
for i in range(S):
|
|
token = sent_tokens[i]
|
|
sent.append( word_vocab.get_index(token) )
|
|
# pad the right end till the max length of the mini batch
|
|
for i in range(S, max_len):
|
|
sent.append( word_vocab.pad_index )
|
|
out_text.append(sent)
|
|
return torch.LongTensor(out_text)
|
|
|
|
def read_labels_tensor(rel_labels, rel_vocab, cuda=False):
|
|
N = len(rel_labels)
|
|
labels_list = []
|
|
for i in range(N):
|
|
token = rel_labels[i]
|
|
labels_list.append( rel_vocab.get_index(token) )
|
|
return torch.LongTensor(labels_list)
|
|
|
|
def read_dataset(datapath, word_vocab, rel_vocab):
|
|
questions = []
|
|
rel_labels = []
|
|
# read questions and label from the datapath - could be train, dev, testls
|
|
with open(datapath) as f:
|
|
for line in f:
|
|
line_items = line.split("\t")
|
|
# add relation
|
|
relation = line_items[1]
|
|
rel_labels.append(relation)
|
|
# add text
|
|
qText = line_items[3]
|
|
tokens = process_tokenize_text(qText)
|
|
questions.append(tokens)
|
|
|
|
dataset = {"word_vocab": word_vocab, "rel_vocab": rel_vocab, "size": len(rel_labels),
|
|
"questions": np.array(questions), "rel_labels": np.array(rel_labels)}
|
|
return dataset
|
|
|