Fix the feature extraction bug (#1942)

Author: w5688414

applications/neural_search/recall/milvus/feature_extract.py

@@ -156,7 +156,13 @@ def predict(self, data, tokenizer):
                 logits = self.output_handle.copy_to_cpu()
                 all_embeddings.append(logits)
                 examples = []
-
+        if (len(examples) > 0):
+            input_ids, segment_ids = batchify_fn(examples)
+            self.input_handles[0].copy_from_cpu(input_ids)
+            self.input_handles[1].copy_from_cpu(segment_ids)
+            self.predictor.run()
+            logits = self.output_handle.copy_to_cpu()
+            all_embeddings.append(logits)
         all_embeddings = np.concatenate(all_embeddings, axis=0)
         np.save('corpus_embedding', all_embeddings)
 

applications/neural_search/recall/milvus/scripts/feature_extract.sh

@@ -1,3 +1,3 @@
 CUDA_VISIBLE_DEVICES=2 python feature_extract.py \
         --model_dir=./output \
-        --corpus_file "data/milvus_data.csv" 
+        --corpus_file "milvus/milvus_data.csv"