From 6039f6471c013d3b08a61b26e925ce80d9f5213e Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Wed, 2 Mar 2022 14:02:26 +0800 Subject: [PATCH] Update Signed-off-by: Jael Gu --- README.md | 21 +++++++++++---------- nlp_longformer.py | 2 +- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/README.md b/README.md index 7ecd753..ec3a213 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # Operator: nlp-longformer -Author: +Author: Kyle He, Jael Gu ## Overview @@ -27,17 +27,18 @@ __call__(self, call_arg_1: xxx) Args: -- call_arg_1: - - xxx(description about call_arg_1) - - supported types: xxx - Returns: +- txt: + - input text in words, sentences, or paragraphs + - supported types: str -The Operator returns a tuple Tuple[('results_1', xxx)] containing following fields: +Returns: -- results_1: - - xxx(description of results_1) - - data type: xxx - - shape: (xxx,) +The Operator returns a tuple Tuple[('feature_vector', numpy.ndarray)] containing following fields: + +- feature_vector: + - the embedding of the text + - data type: numpy.ndarray + - shape: (x, dim) where x is number of vectors and dim is dimension of vector depending on model_name ## Requirements diff --git a/nlp_longformer.py b/nlp_longformer.py index 955060c..2d7c1bf 100644 --- a/nlp_longformer.py +++ b/nlp_longformer.py @@ -27,7 +27,7 @@ class NlpLongformer(NNOperator): input_ids = torch.tensor(self.tokenizer.encode(txt)).unsqueeze(0) attention_mask = None outs = self.model(input_ids, attention_mask=attention_mask, output_hidden_states=True) - feature_vector = outs[1].squeeze() + feature_vector = outs[1].squeeze(0) Outputs = NamedTuple('Outputs', [('feature_vector', numpy.ndarray)]) return Outputs(feature_vector.detach().numpy())