enhance predict dataset

Author: kinghuin

paddlehub/dataset/base_nlp_dataset.py

@@ -21,6 +21,7 @@
 import csv
 
 from paddlehub.dataset import InputExample, BaseDataset
+from paddlehub.common.logger import logger
 
 
 class BaseNLPDataset(BaseDataset):
@@ -51,6 +52,7 @@ def __init__(self,
 
     def _read_file(self, input_file, phase=None):
         """Reads a tab separated value file."""
+        has_warned = False
         with io.open(input_file, "r", encoding="UTF-8") as file:
             reader = csv.reader(file, delimiter="\t", quotechar=None)
             examples = []
@@ -59,22 +61,38 @@ def _read_file(self, input_file, phase=None):
                     ncol = len(line)
                     if self.if_file_with_header[phase]:
                         continue
-                if ncol == 1:
-                    if phase == "predict":
-                        example = InputExample(guid=i, text_a=line[0])
-                    else:
+                if phase != "predict":
+                    if ncol == 1:
                         raise Exception(
                             "the %s file: %s only has one column but it is not a predict file"
                             % (phase, input_file))
-                elif ncol == 2:
-                    example = InputExample(
-                        guid=i, text_a=line[0], label=line[1])
-                elif ncol == 3:
-                    example = InputExample(
-                        guid=i, text_a=line[0], text_b=line[1], label=line[2])
+                    elif ncol == 2:
+                        example = InputExample(
+                            guid=i, text_a=line[0], label=line[1])
+                    elif ncol == 3:
+                        example = InputExample(
+                            guid=i,
+                            text_a=line[0],
+                            text_b=line[1],
+                            label=line[2])
+                    else:
+                        raise Exception(
+                            "the %s file: %s has too many columns (should <=3)"
+                            % (phase, input_file))
                 else:
-                    raise Exception(
-                        "the %s file: %s has too many columns (should <=3)" %
-                        (phase, input_file))
+                    if ncol == 1:
+                        example = InputExample(guid=i, text_a=line[0])
+                    elif ncol == 2:
+                        if not has_warned:
+                            logger.warning(
+                                "the predict file: %s has 2 columns, as it is a predict file, the second one will be regarded as text_b"
+                                % (input_file))
+                            has_warned = True
+                        example = InputExample(
+                            guid=i, text_a=line[0], text_b=line[1])
+                    else:
+                        raise Exception(
+                            "the predict file: %s has too many columns (should <=2)"
+                            % (input_file))
                 examples.append(example)
             return examples