add structure tasks: sudoku & text2sql

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11314581

modelscope/metainfo.py

@@ -346,9 +346,11 @@ class Pipelines(object):
     image_text_retrieval = 'image-text-retrieval'
     ofa_ocr_recognition = 'ofa-ocr-recognition'
     ofa_asr = 'ofa-asr'
-    document_vl_embedding = 'document-vl-embedding'
+    ofa_sudoku = 'ofa-sudoku'
+    ofa_text2sql = 'ofa-text2sql'
     video_captioning = 'video-captioning'
     video_question_answering = 'video-question-answering'
+    document_vl_embedding = 'document-vl-embedding'
 
     # science tasks
     protein_structure = 'unifold-protein-structure'

modelscope/models/multi_modal/ofa/configuration_mmspeech.py

@@ -209,6 +209,12 @@ class MMSpeechConfig(PretrainedConfig):
                  use_ofasys=False,
                  vit_type='vit_base',
                  vit_drop_path_rate=0.0,
+                 use_gamma_feature=False,
+                 gamma=1.0,
+                 exclude_mlp=True,
+                 temperature_init_value=None,
+                 remove_decoder_type_embedding=False,
+                 mlp_dim=512,
                  required_seq_len_multiple=2,
                  encoder_pos_conv_depth=5,
                  encoder_conv_pos=95,
@@ -279,6 +285,15 @@ class MMSpeechConfig(PretrainedConfig):
         self.use_ofasys = use_ofasys
         self.vit_type = vit_type
         self.vit_drop_path_rate = vit_drop_path_rate
+        self.use_gamma_feature = use_gamma_feature
+
+        # add some new features from ofa
+        self.use_gamma_feature = use_gamma_feature
+        self.gamma = gamma
+        self.exclude_mlp = exclude_mlp
+        self.temperature_init_value = temperature_init_value
+        self.remove_decoder_type_embedding = remove_decoder_type_embedding
+        self.mlp_dim = mlp_dim
 
         # FP16 optimization
         self.required_seq_len_multiple = required_seq_len_multiple

modelscope/models/multi_modal/ofa/configuration_ofa.py

@@ -216,6 +216,12 @@ class OFAConfig(PretrainedConfig):
                  use_ofasys=False,
                  vit_type='vit_base',
                  vit_drop_path_rate=0.0,
+                 use_gamma_feature=False,
+                 gamma=1.0,
+                 exclude_mlp=True,
+                 temperature_init_value=None,
+                 remove_decoder_type_embedding=False,
+                 mlp_dim=512,
                  **kwargs):
         self.vocab_size = vocab_size
         self.max_position_embeddings = max_position_embeddings
@@ -265,6 +271,14 @@ class OFAConfig(PretrainedConfig):
         self.vit_type = vit_type
         self.vit_drop_path_rate = vit_drop_path_rate
 
+        # add some new features from ofa
+        self.use_gamma_feature = use_gamma_feature
+        self.gamma = gamma
+        self.exclude_mlp = exclude_mlp
+        self.temperature_init_value = temperature_init_value
+        self.remove_decoder_type_embedding = remove_decoder_type_embedding
+        self.mlp_dim = mlp_dim
+
         super().__init__(
             pad_token_id=pad_token_id,
             bos_token_id=bos_token_id,

modelscope/models/multi_modal/ofa/modeling_mmspeech.py

@@ -11,12 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-""" PyTorch OFA model."""
+""" PyTorch OFA-MMSpeech model."""
 
 import math
-import random
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple
+from typing import Optional, Tuple
 
 import numpy as np
 import torch
@@ -27,22 +26,17 @@ from fairseq.modules import LayerNorm, SamePad, TransposeLast
 from fairseq.modules.transformer_sentence_encoder import init_bert_params
 from fairseq.utils import index_put
 from packaging import version
-from torch import Tensor, nn
+from torch import nn
 from torch.nn import functional as F
-from transformers.activations import ACT2FN
 from transformers.file_utils import (ModelOutput, add_code_sample_docstrings,
                                      add_start_docstrings,
                                      add_start_docstrings_to_model_forward)
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput,
-    Seq2SeqModelOutput)
-from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 
 from .configuration_mmspeech import MMSpeechConfig
 from .generate import utils
 from .modeling_ofa import (Embedding, OFADecoder, OFAModel, OFAPreTrainedModel,
-                           _expand_mask, shift_tokens_right)
+                           _expand_mask)
 
 logger = logging.get_logger()
 

modelscope/models/multi_modal/ofa/modeling_ofa.py

@@ -461,6 +461,18 @@ class OFAEncoderLayer(nn.Module):
         self.drop_path = DropPath(
             drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
 
+        self.use_gamma_feature = config.use_gamma_feature
+        if self.use_gamma_feature:
+            gamma = getattr(config, 'gamma', 1.)
+
+            # `OFA.from_pretrain()` method will replace the `gamma` to `weight`
+            # in the model key. Here, change the parameters like `xxx_gamma_xxx`
+            # to `xxx_weight_xxx` to adapt this transformation.
+            self.weight_self_attn = nn.Parameter(
+                torch.ones(self.embed_dim) * gamma, requires_grad=True)
+            self.weight_ffn = nn.Parameter(
+                torch.ones(self.embed_dim) * gamma, requires_grad=True)
+
     def residual_connection(self, x, residual):
         r"""
         Residual connection with drop path.
@@ -499,6 +511,8 @@ class OFAEncoderLayer(nn.Module):
         if self.self_attn_mid_layer_norm:
             hidden_states = self.self_attn_mid_layer_norm(hidden_states)
         hidden_states = self.dropout(hidden_states)
+        if self.use_gamma_feature:
+            hidden_states = self.weight_self_attn * hidden_states
         hidden_states = self.residual_connection(hidden_states, residual)
         if not self.normalize_before:
             hidden_states = self.self_attn_layer_norm(hidden_states)
@@ -513,6 +527,8 @@ class OFAEncoderLayer(nn.Module):
             hidden_states = self.ffn_layer_norm(hidden_states)
         hidden_states = self.fc2(hidden_states)
         hidden_states = self.dropout(hidden_states)
+        if self.use_gamma_feature:
+            hidden_states = self.weight_ffn * hidden_states
         hidden_states = self.residual_connection(hidden_states, residual)
         if not self.normalize_before:
             hidden_states = self.final_layer_norm(hidden_states)
@@ -578,6 +594,20 @@ class OFADecoderLayer(nn.Module):
         self.drop_path = DropPath(
             drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
 
+        self.use_gamma_feature = config.use_gamma_feature
+        if self.use_gamma_feature:
+            gamma = getattr(config, 'gamma', 1.)
+
+            # `OFA.from_pretrain()` method will replace the `gamma` to `weight`
+            # in the model key. Here, change the parameters like `xxx_gamma_xxx`
+            # to `xxx_weight_xxx` to adapt this transformation.
+            self.weight_self_attn = nn.Parameter(
+                torch.ones(self.embed_dim) * gamma, requires_grad=True)
+            self.weight_cross_attn = nn.Parameter(
+                torch.ones(self.embed_dim) * gamma, requires_grad=True)
+            self.weight_ffn = nn.Parameter(
+                torch.ones(self.embed_dim) * gamma, requires_grad=True)
+
     def residual_connection(self, x, residual):
         r"""
         Residual connection with drop path.
@@ -629,6 +659,8 @@ class OFADecoderLayer(nn.Module):
         if self.self_attn_mid_layer_norm:
             hidden_states = self.self_attn_mid_layer_norm(hidden_states)
         hidden_states = self.dropout(hidden_states)
+        if self.use_gamma_feature:
+            hidden_states = self.weight_self_attn * hidden_states
         hidden_states = self.residual_connection(hidden_states, residual)
         if not self.normalize_before:
             hidden_states = self.self_attn_layer_norm(hidden_states)
@@ -654,6 +686,8 @@ class OFADecoderLayer(nn.Module):
             if self.cross_attn_mid_layer_norm:
                 hidden_states = self.cross_attn_mid_layer_norm(hidden_states)
             hidden_states = self.dropout(hidden_states)
+            if self.use_gamma_feature:
+                hidden_states = self.weight_cross_attn * hidden_states
             hidden_states = self.residual_connection(hidden_states, residual)
             if not self.normalize_before:
                 hidden_states = self.cross_attn_layer_norm(hidden_states)
@@ -671,6 +705,8 @@ class OFADecoderLayer(nn.Module):
             hidden_states = self.ffn_layer_norm(hidden_states)
         hidden_states = self.fc2(hidden_states)
         hidden_states = self.dropout(hidden_states)
+        if self.use_gamma_feature:
+            hidden_states = self.weight_ffn * hidden_states
         hidden_states = self.residual_connection(hidden_states, residual)
         if not self.normalize_before:
             hidden_states = self.final_layer_norm(hidden_states)
@@ -1961,6 +1997,14 @@ class OFAModel(OFAPreTrainedModel):
         self.decoder = OFADecoder(config, shared)
         self.use_ofasys = config.use_ofasys
 
+        # exclude mlp head as default
+        if not getattr(config, 'exclude_mlp', True):
+            self.mlp_head = Linear(config.d_model, config.mlp_dim)
+        # None temperature_init_value as default
+        if config.temperature_init_value:
+            self.temp = nn.Parameter(config.temperature_init_value
+                                     * torch.ones([]))
+
         # Initialize weights and apply final processing
         self.post_init()
 

modelscope/models/multi_modal/ofa/utils/constant.py

@@ -11,5 +11,7 @@ OFA_TASK_KEY_MAPPING = {
     Tasks.text_classification: OutputKeys.LABELS,
     Tasks.image_classification: OutputKeys.LABELS,
     Tasks.visual_entailment: OutputKeys.LABELS,
-    Tasks.auto_speech_recognition: OutputKeys.TEXT
+    Tasks.auto_speech_recognition: OutputKeys.TEXT,
+    Tasks.sudoku: OutputKeys.TEXT,
+    Tasks.text2sql: OutputKeys.TEXT,
 }

modelscope/models/multi_modal/ofa_for_all_tasks.py

@@ -38,6 +38,8 @@ __all__ = ['OfaForAllTasks']
 @MODELS.register_module(Tasks.text_summarization, module_name=Models.ofa)
 @MODELS.register_module(Tasks.text_classification, module_name=Models.ofa)
 @MODELS.register_module(Tasks.auto_speech_recognition, module_name=Models.ofa)
+@MODELS.register_module(Tasks.sudoku, module_name=Models.ofa)
+@MODELS.register_module(Tasks.text2sql, module_name=Models.ofa)
 class OfaForAllTasks(TorchModel):
     r"""
     All ofa tasks using uniform ofa model structure. So far, we support three types of tasks:
@@ -94,7 +96,7 @@ class OfaForAllTasks(TorchModel):
         self.cfg = Config.from_file(
             osp.join(model_dir, ModelFile.CONFIGURATION))
         multimodal_type = self.cfg.model.get('multimodal_type', 'default')
-        if multimodal_type == 'default':
+        if multimodal_type in ['default', 'text2sql']:
             model = OFAModel.from_pretrained(model_dir)
         elif multimodal_type == 'mmspeech':
             model = MMSpeechModel.from_pretrained(model_dir)
@@ -123,6 +125,13 @@ class OfaForAllTasks(TorchModel):
                 self.tokenizer.add_tokens(
                     ['<audio_{}>'.format(i) for i in range(30000)])
                 self.cfg.update({'num_bins': 0, 'num_codes': 30000})
+            elif multimodal_type == 'text2sql':
+                self.tokenizer.add_tokens(
+                    ['<code_{}>'.format(i) for i in range(8192)])
+                self.tokenizer.add_tokens(
+                    ['<bin_{}>'.format(i) for i in range(1000)])
+                self.cfg.update({'num_bins': 1000, 'num_codes': 8192})
+                self.tokenizer.add_tokens(['>=', '<='])
 
         self.batch_size = self.cfg.model.get('batch_size', 1)
         self.patch_image_size = self.cfg.model.get('patch_image_size', 480)
@@ -177,6 +186,8 @@ class OfaForAllTasks(TorchModel):
             Tasks.text_classification: inference_d[self.gen_type],
             Tasks.image_classification: inference_d[self.gen_type],
             Tasks.auto_speech_recognition: self._text_gen_inference,
+            Tasks.sudoku: self._text_gen_inference,
+            Tasks.text2sql: self._text_gen_inference,
         }
         pattern_str = '((?<=[^ a-zA-Z0-9.,:!?]) +| +(?=[^ a-zA-Z0-9.,:!?]))'
         self.pattern = re.compile(pattern_str)

modelscope/outputs/outputs.py

@@ -73,6 +73,8 @@ TASK_OUTPUTS = {
     #    "text": "电子元器件提供BOM配单"
     # }
     Tasks.ocr_recognition: [OutputKeys.TEXT],
+    Tasks.sudoku: [OutputKeys.TEXT],
+    Tasks.text2sql: [OutputKeys.TEXT],
 
     # document vl embedding for single sample
     # {

modelscope/pipeline_inputs.py

@@ -201,6 +201,12 @@ TASK_INPUTS = {
         'src': InputType.LIST,
         'ref': InputType.LIST,
     },
+    Tasks.sudoku:
+    InputType.TEXT,
+    Tasks.text2sql: {
+        'text': InputType.TEXT,
+        'database': InputType.TEXT
+    },
 
     # ============ audio tasks ===================
     Tasks.auto_speech_recognition:

modelscope/pipelines/multi_modal/sudoku_pipeline.py

@@ -0,0 +1,53 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, Optional, Union
+
+import torch
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.multi_modal import OfaForAllTasks
+from modelscope.pipelines.base import Model, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.pipelines.util import batch_process
+from modelscope.preprocessors import OfaPreprocessor, Preprocessor
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(Tasks.sudoku, module_name=Pipelines.ofa_sudoku)
+class SudokuPipeline(Pipeline):
+    R"""
+    pipeline for sudoku solving
+    """
+
+    def __init__(self,
+                 model: Union[Model, str],
+                 preprocessor: Optional[Preprocessor] = None,
+                 **kwargs):
+        """
+        use `model` and `preprocessor` to create a pipeline for solving sudoku
+        Args:
+            model: model id on modelscope hub.
+        """
+        super().__init__(model=model, preprocessor=preprocessor, **kwargs)
+        self.model.eval()
+        if preprocessor is None:
+            if isinstance(self.model, OfaForAllTasks):
+                self.preprocessor = OfaPreprocessor(self.model.model_dir)
+            else:
+                raise 'no preprocessor is provided'
+
+    def _batch(self, data):
+        if isinstance(self.model, OfaForAllTasks):
+            return batch_process(self.model, data)
+        else:
+            return super(SudokuPipeline, self)._batch(data)
+
+    def forward(self, inputs: Dict[str, Any],
+                **forward_params) -> Dict[str, Any]:
+        with torch.no_grad():
+            return super().forward(inputs, **forward_params)
+
+    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        return inputs

modelscope/pipelines/multi_modal/text2sql_pipeline.py

@@ -0,0 +1,51 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict, Optional, Union
+
+import torch
+
+from modelscope.metainfo import Pipelines
+from modelscope.models.multi_modal import OfaForAllTasks
+from modelscope.pipelines.base import Model, Pipeline
+from modelscope.pipelines.builder import PIPELINES
+from modelscope.pipelines.util import batch_process
+from modelscope.preprocessors import OfaPreprocessor, Preprocessor
+from modelscope.utils.constant import Tasks
+from modelscope.utils.logger import get_logger
+
+logger = get_logger()
+
+
+@PIPELINES.register_module(Tasks.text2sql, module_name=Pipelines.ofa_text2sql)
+class TextToSqlPipeline(Pipeline):
+    R"""
+    pipeline for text to sql task
+    """
+
+    def __init__(self,
+                 model: Union[Model, str],
+                 preprocessor: Optional[Preprocessor] = None,
+                 **kwargs):
+        """
+        use `model` and `preprocessor` to create a pipeline for text2sql task
+        Args:
+            model: model id on modelscope hub.
+        """
+        super().__init__(model=model, preprocessor=preprocessor, **kwargs)
+        self.model.eval()
+        if preprocessor is None:
+            if isinstance(self.model, OfaForAllTasks):
+                self.preprocessor = OfaPreprocessor(self.model.model_dir)
+
+    def _batch(self, data):
+        if isinstance(self.model, OfaForAllTasks):
+            return batch_process(self.model, data)
+        else:
+            return super(TextToSqlPipeline, self)._batch(data)
+
+    def forward(self, inputs: Dict[str, Any],
+                **forward_params) -> Dict[str, Any]:
+        with torch.no_grad():
+            return super().forward(inputs, **forward_params)
+
+    def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
+        return inputs

modelscope/preprocessors/multi_modal.py

@@ -56,7 +56,9 @@ class OfaPreprocessor(Preprocessor):
             Tasks.text_classification: OfaTextClassificationPreprocessor,
             Tasks.text_summarization: OfaSummarizationPreprocessor,
             Tasks.text_to_image_synthesis: OfaTextToImageSynthesisPreprocessor,
-            Tasks.auto_speech_recognition: OfaASRPreprocessor
+            Tasks.auto_speech_recognition: OfaASRPreprocessor,
+            Tasks.sudoku: OfaSudokuPreprocessor,
+            Tasks.text2sql: OfaTextToSqlPreprocessor
         }
         model_dir = model_dir if osp.exists(model_dir) else snapshot_download(
             model_dir, user_agent={Invoke.KEY: Invoke.PREPROCESSOR})

modelscope/preprocessors/ofa/__init__.py

@@ -3,7 +3,9 @@ from .asr import OfaASRPreprocessor
 from .image_captioning import OfaImageCaptioningPreprocessor
 from .image_classification import OfaImageClassificationPreprocessor
 from .ocr_recognition import OfaOcrRecognitionPreprocessor
+from .sudoku import OfaSudokuPreprocessor
 from .summarization import OfaSummarizationPreprocessor
+from .text2sql import OfaTextToSqlPreprocessor
 from .text_classification import OfaTextClassificationPreprocessor
 from .text_to_image_synthesis import OfaTextToImageSynthesisPreprocessor
 from .visual_entailment import OfaVisualEntailmentPreprocessor

modelscope/preprocessors/ofa/base.py

@@ -48,6 +48,8 @@ class OfaBasePreprocessor:
         # there will be no need to use param: use_bpe
         tokenizer.add_tokens(['<code_{}>'.format(i) for i in range(8192)])
         tokenizer.add_tokens(['<bin_{}>'.format(i) for i in range(1000)])
+        if self.cfg.model.get('multimodal_type', 'default') == 'text2sql':
+            tokenizer.add_tokens(['>=', '<='])
         self.tokenizer = tokenizer
         self.bos_item = torch.LongTensor([tokenizer.bos_token_id])
         self.pad_item = torch.LongTensor([tokenizer.pad_token_id])

modelscope/preprocessors/ofa/sudoku.py

@@ -0,0 +1,110 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+from typing import Any, Dict
+
+import numpy as np
+import torch
+
+from modelscope.utils.constant import ModeKeys
+from .base import OfaBasePreprocessor
+
+
+class OfaSudokuPreprocessor(OfaBasePreprocessor):
+    r"""
+    OFA preprocessor for sudoku tasks
+    """
+
+    def __init__(self,
+                 cfg,
+                 model_dir,
+                 mode=ModeKeys.INFERENCE,
+                 *args,
+                 **kwargs):
+        """preprocess the data
+
+        Args:
+            cfg(modelscope.utils.config.ConfigDict) : model config
+            model_dir (str): model path,
+            mode: preprocessor mode (model mode)
+        """
+        super(OfaSudokuPreprocessor, self).__init__(cfg, model_dir, mode,
+                                                    *args, **kwargs)
+
+        self.instruction_text = self.cfg.model.get('prompt',
+                                                   ' solve the sudoku .')
+        self.seg_embedding = self.cfg.get('seg_embedding', False)
+        self.max_struct_length = self.cfg.get('max_struct_length', 256)
+        if self.seg_embedding:
+            self.input_puzzle_row = []
+            self.input_puzzle_col = []
+            for idx in range(9):
+                for jdx in range(9):
+                    self.input_puzzle_row.append(jdx + 1)
+                    self.input_puzzle_col.append(idx + 1)
+                    if not (idx == 8 and jdx == 8):
+                        self.input_puzzle_row.append(0)
+                        self.input_puzzle_col.append(0)
+            self.input_puzzle_col = torch.tensor(self.input_puzzle_col)
+            self.input_puzzle_row = torch.tensor(self.input_puzzle_row)
+
+            instruct_seg = torch.zeros_like(
+                self.tokenize_text(self.instruction_text))
+            input_puzzle_col = torch.cat([self.input_puzzle_col, instruct_seg])
+            input_puzzle_row = torch.cat([self.input_puzzle_row, instruct_seg])
+            self.input_puzzle_col = torch.cat(
+                [self.bos_item, input_puzzle_col, self.eos_item])
+            self.input_puzzle_row = torch.cat(
+                [self.bos_item, input_puzzle_row, self.eos_item])
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if self.mode == ModeKeys.TRAIN:
+            return self._build_train_sample(data)
+        else:
+            return self._build_infer_sample(data)
+
+    def _build_train_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        build sample for training tasks.
+
+        step 1. execute the `_build_infer_sample` function to get a batch sample
+            for inference.
+        step 2. process the label data for training.
+        """
+        sample = self._build_infer_sample(data)
+        target = sample['label']
+        target_token_list = target.lower().strip().split()
+        target = ' '.join(target_token_list[:self.max_tgt_length])
+        sample['target'] = self.tokenize_text(target, add_bos=False)
+        sample['prev_output_tokens'] = torch.cat(
+            [self.bos_item, sample['target'][:-1]])
+        return sample
+
+    def _build_infer_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        build sample for inference tasks.
+
+        step 1. Get the input random masked sudoku text input, which shold be
+            generated like below pseudo code.
+            >>> sudo = np.random.randint(1, 9, size=(9, 9)) # a pseudo sudoku
+            >>> sudo_text = " | ".join(" : ".join(str(c) for c in row) \
+            >>>             for row in sudo)
+        step 2. Limit the length, tokenize the input text and add the bos token
+            to the front of the input as source input.
+        step 3. Add a pseodo ids for every input.
+        """
+        assert 'text' in self.column_map and 'text' in data, \
+            'there must be `text` column in task key map and source data'
+        text = data[self.column_map['text']]  # equal data['text']
+        text = ' '.join(text.lower().strip().split()[:self.max_struct_length])
+        src_item = self.tokenize_text(text + self.instruction_text)
+        src_item = src_item[:(self.max_src_length + self.max_struct_length)]
+
+        sample = {'id': 0.0, 'source': src_item}
+
+        if self.seg_embedding:
+            sample['seg_row_tokens'] = self.input_puzzle_row
+            sample['seg_col_tokens'] = self.input_puzzle_col
+
+        if 'solution' in self.column_map and self.column_map[
+                'solution'] in data:
+            sample['label'] = ' {}'.format(data[self.column_map['solution']])
+        return sample

modelscope/preprocessors/ofa/text2sql.py

@@ -0,0 +1,446 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import os
+import random
+import re
+from typing import Any, Dict, List
+
+import torch
+
+from modelscope.utils.constant import ModeKeys
+from .base import OfaBasePreprocessor
+from .utils.bridge_content_encoder import get_database_matches
+from .utils.get_tables import dump_db_json_schema
+
+
+class OfaTextToSqlPreprocessor(OfaBasePreprocessor):
+    r"""
+    OFA preprocessor for text to sql tasks
+    """
+
+    def __init__(self,
+                 cfg,
+                 model_dir,
+                 mode=ModeKeys.INFERENCE,
+                 *args,
+                 **kwargs):
+        """preprocess the data
+
+        Args:
+            cfg(modelscope.utils.config.ConfigDict) : model config
+            model_dir (str): model path,
+            mode: preprocessor mode (model mode)
+        """
+        super(OfaTextToSqlPreprocessor, self).__init__(cfg, model_dir, mode,
+                                                       *args, **kwargs)
+
+        self.instruction_text = self.cfg.model.get('prompt',
+                                                   ' . generating sql code.')
+        self.max_struct_length = self.cfg.get('max_struct_length', 256)
+        self.separator = '\t'
+        self.db_schema_cache = {}
+        self.database_path = os.path.join(
+            os.path.abspath(model_dir), 'database')
+
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if self.mode == ModeKeys.TRAIN:
+            return self._build_train_sample(data)
+        else:
+            return self._build_infer_sample(data)
+
+    def _build_train_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        build sample for training tasks.
+
+        step 1. Get the input question and database id from text input
+        step 2. Get the database structure input
+        step 3. Add a pseudo ids for every input.
+        step 4. Calculate the target and previous output items.
+        """
+        assert 'text' in self.column_map and 'text' in data, \
+            'there must be `text` column in task key map and source data'
+        text = data[self.column_map['text']]  # equal data['text']
+        texts = text.split(self.separator)
+        assert len(
+            texts
+        ) == 3, 'invalid input, should contain query, question and database id'
+        query, question, db_id = texts
+
+        # construct struct input
+        if db_id not in self.db_schema_cache:
+            self.db_schema_cache[db_id] = dump_db_json_schema(
+                self.database_path + '/' + db_id + '/' + db_id + '.sqlite',
+                db_id)
+
+        question = ' '.join(question.strip().split()[:self.max_src_length])
+
+        seq_inputs = seq2seq_input(query, question, db_id, self.database_path,
+                                   self.db_schema_cache[db_id], self.cfg.model,
+                                   True)
+        struct_in = seq_inputs['struct_in']
+        text = seq_inputs['text_in']
+        seq_out = seq_inputs['seq_out']
+        db_struct = seq_inputs['db_struct']
+
+        text = '{} ; structured knowledge: {}'.format(
+            text, struct_in) + self.instruction_text
+        src_item = self.tokenize_text(text + self.instruction_text)
+        src_item = src_item[:(self.max_src_length + self.max_struct_length
+                              + 20)]
+
+        tgt_item = self.tokenize_text(
+            ' {}'.format(seq_out), add_bos=False,
+            add_eos=False)[:self.max_tgt_length]
+        target_item = torch.cat([tgt_item, self.eos_item])
+        prev_output_item = torch.cat([self.bos_item, tgt_item])
+
+        sample = {
+            'id': 0.0,
+            'source': src_item,
+            'target': target_item,
+            'prev_output_tokens': prev_output_item,
+            'db_struct': db_struct
+        }
+
+        return sample
+
+    def _build_infer_sample(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        r"""
+        build sample for inference tasks.
+
+        step 1. Get the input question and database id from text input
+        step 2. Get the database structure input
+        step 3. Add a pseudo ids for every input.
+        """
+        assert 'text' in self.column_map and 'text' in data, \
+            'there must be `text` column in task key map and source data'
+        text = data[self.column_map['text']]  # equal data['text']
+        db_id = data.get(self.column_map['database'], 'culture_company')
+        db_id = db_id.strip()
+
+        # construct struct input
+        if db_id not in self.db_schema_cache:
+            self.db_schema_cache[db_id] = dump_db_json_schema(
+                self.database_path + '/' + db_id + '/' + db_id + '.sqlite',
+                db_id)
+
+        text = ' '.join(text.strip().split()[:self.max_src_length])
+
+        seq_inputs = seq2seq_input(None, text, db_id, self.database_path,
+                                   self.db_schema_cache[db_id], self.cfg.model)
+        struct_in = seq_inputs['struct_in']
+        db_struct = seq_inputs['db_struct']
+        text = '{} ; structured knowledge: {}'.format(
+            text, struct_in) + self.instruction_text
+        src_item = self.tokenize_text(text + self.instruction_text)
+        src_item = src_item[:(self.max_src_length + self.max_struct_length
+                              + 20)]
+
+        sample = {'id': 0.0, 'source': src_item, 'db_struct': db_struct}
+
+        if 'solution' in self.column_map and self.column_map[
+                'solution'] in data:
+            sample['label'] = ' {}'.format(data[self.column_map['solution']])
+        return sample
+
+
+def seq2seq_input(query,
+                  question,
+                  db_id,
+                  db_path,
+                  schema,
+                  args,
+                  is_train=False):
+    ex = form_input_for_construction(query, question, db_id, db_path, schema)
+    serialized_schema = spider_add_serialized_schema(
+        ex, args)['serialized_schema'].strip()
+    if not is_train:
+        return {
+            'struct_in': serialized_schema,
+            'text_in': question,
+            'db_struct': ex
+        }
+    question, seq_out = spider_pre_process_one_function(ex, args)
+    return {
+        'struct_in': serialized_schema,
+        'text_in': question,
+        'seq_out': seq_out,
+        'db_struct': ex
+    }
+
+
+def spider_pre_process_one_function(item: dict, args):
+    prefix = ''
+
+    seq_out = spider_get_target(
+        query=item['query'],
+        db_id=item['db_id'],
+        normalize_query=True,
+        target_with_db_id=args.target_with_db_id,
+    )
+
+    return prefix + item['question'].strip(), seq_out
+
+
+def spider_get_target(
+    query: str,
+    db_id: str,
+    normalize_query: bool,
+    target_with_db_id: bool,
+) -> str:
+    _normalize = normalize if normalize_query else (lambda x: x)
+    return f'{db_id} | {_normalize(query)}' if target_with_db_id else _normalize(
+        query)
+
+
+def normalize(query: str) -> str:
+
+    def comma_fix(s):
+        # Remove spaces in front of commas
+        return s.replace(' , ', ', ')
+
+    def white_space_fix(s):
+        # Remove double and triple spaces
+        return ' '.join(s.split())
+
+    def lower(s):
+        # Convert everything except text between (single or double) quotation marks to lower case
+        return re.sub(r"\b(?<!['\"])(\w+)(?!['\"])\b",
+                      lambda match: match.group(1).lower(), s)
+
+    return comma_fix(white_space_fix(lower(query)))
+
+
+def spider_add_serialized_schema(ex: dict, args) -> dict:
+    if getattr(args, 'schema_serialization_with_nl'):
+        serialized_schema = serialize_schema_natural_language(
+            question=ex['question'],
+            db_path=ex['db_path'],
+            db_id=ex['db_id'],
+            db_column_names=ex['db_column_names'],
+            db_table_names=ex['db_table_names'],
+            db_primary_keys=ex['db_primary_keys'],
+            db_foreign_keys=ex['db_foreign_keys'],
+            schema_serialization_with_db_content=args.
+            schema_serialization_with_db_content,
+            normalize_query=True,
+        )
+    else:
+        serialized_schema = serialize_schema(
+            question=ex['question'],
+            db_path=ex['db_path'],
+            db_id=ex['db_id'],
+            db_column_names=ex['db_column_names'],
+            db_table_names=ex['db_table_names'],
+            schema_serialization_type='peteshaw',
+            schema_serialization_randomized=False,
+            schema_serialization_with_db_id=True,
+            schema_serialization_with_db_content=args.
+            schema_serialization_with_db_content,
+            normalize_query=True,
+        )
+    return {'serialized_schema': serialized_schema}
+
+
+def serialize_schema_natural_language(
+    question: str,
+    db_path: str,
+    db_id: str,
+    db_column_names: Dict[str, str],
+    db_table_names: List[str],
+    db_primary_keys,
+    db_foreign_keys,
+    schema_serialization_with_db_content: bool = False,
+    normalize_query: bool = True,
+) -> str:
+    overall_description = f'{db_id} contains tables such as ' \
+                          f'{", ".join([name.lower() if normalize_query else name for name in db_table_names])}.'
+
+    def table_description_primary_key_template(primary_key):
+        return f'{primary_key} is the primary key.'
+
+    def table_description(name, column_names):
+        return f'Table {name} has columns such as {", ".join(column_names)}.'
+
+    def value_description(cv_pairs):
+        return f'{"".join(["The {} contains values such as {}.".format(column, value) for column, value in cv_pairs])}'
+
+    def foreign_key_description(table_1, column_1, table_2, column_2):
+        return f'The {column_1} of {table_1} is the foreign key of {column_2} of {table_2}.'
+
+    db_primary_keys = db_primary_keys['column_id']
+    db_foreign_keys = list(
+        zip(db_foreign_keys['column_id'], db_foreign_keys['other_column_id']))
+
+    descriptions = [overall_description]
+    db_table_name_strs = []
+    db_column_name_strs = []
+    value_sep = ', '
+    for table_id, table_name in enumerate(db_table_names):
+        table_name_str = table_name.lower() if normalize_query else table_name
+        db_table_name_strs.append(table_name_str)
+        columns = []
+        column_value_pairs = []
+        primary_keys = []
+        for column_id, (x, y) in enumerate(
+                zip(db_column_names['table_id'],
+                    db_column_names['column_name'])):
+            if column_id == 0:
+                continue
+            column_str = y.lower() if normalize_query else y
+            db_column_name_strs.append(column_str)
+            if x == table_id:
+                columns.append(column_str)
+                if column_id in db_primary_keys:
+                    primary_keys.append(column_str)
+                if schema_serialization_with_db_content:
+                    matches = get_database_matches(
+                        question=question,
+                        table_name=table_name,
+                        column_name=y,
+                        db_path=(db_path + '/' + db_id + '/' + db_id
+                                 + '.sqlite'),
+                    )
+                    if matches:
+                        column_value_pairs.append(
+                            (column_str, value_sep.join(matches)))
+
+        table_description_columns_str = table_description(
+            table_name_str, columns)
+        descriptions.append(table_description_columns_str)
+        table_description_primary_key_str = table_description_primary_key_template(
+            ', '.join(primary_keys))
+        descriptions.append(table_description_primary_key_str)
+        if len(column_value_pairs) > 0:
+            value_description_str = value_description(column_value_pairs)
+            descriptions.append(value_description_str)
+
+    for x, y in db_foreign_keys:
+        # get the table and column of x
+        x_table_name = db_table_name_strs[db_column_names['table_id'][x]]
+        x_column_name = db_column_name_strs[x]
+        # get the table and column of y
+        y_table_name = db_table_name_strs[db_column_names['table_id'][y]]
+        y_column_name = db_column_name_strs[y]
+        foreign_key_description_str = foreign_key_description(
+            x_table_name, x_column_name, y_table_name, y_column_name)
+        descriptions.append(foreign_key_description_str)
+    return ' '.join(descriptions)
+
+
+def serialize_schema(
+    question: str,
+    db_path: str,
+    db_id: str,
+    db_column_names: Dict[str, str],
+    db_table_names: List[str],
+    schema_serialization_type: str = 'peteshaw',
+    schema_serialization_randomized: bool = False,
+    schema_serialization_with_db_id: bool = True,
+    schema_serialization_with_db_content: bool = False,
+    normalize_query: bool = True,
+) -> str:
+    if schema_serialization_type == 'verbose':
+        db_id_str = 'Database: {db_id}. '
+        table_sep = '. '
+        table_str = 'Table: {table}. Columns: {columns}'
+        column_sep = ', '
+        column_str_with_values = '{column} ({values})'
+        column_str_without_values = '{column}'
+        value_sep = ', '
+    elif schema_serialization_type == 'peteshaw':
+        # see https://github.com/google-research/language/blob/master/language/nqg/tasks/spider/append_schema.py#L42
+        db_id_str = ' | {db_id}'
+        table_sep = ''
+        table_str = ' | {table} : {columns}'
+        column_sep = ' , '
+        column_str_with_values = '{column} ( {values} )'
+        column_str_without_values = '{column}'
+        value_sep = ' , '
+    else:
+        raise NotImplementedError
+
+    def get_column_str(table_name: str, column_name: str) -> str:
+        column_name_str = column_name.lower(
+        ) if normalize_query else column_name
+        if schema_serialization_with_db_content:
+            # print("testing")
+            matches = get_database_matches(
+                question=question,
+                table_name=table_name,
+                column_name=column_name,
+                db_path=(db_path + '/' + db_id + '/' + db_id + '.sqlite'),
+            )
+            if matches:
+                return column_str_with_values.format(
+                    column=column_name_str, values=value_sep.join(matches))
+            else:
+                return column_str_without_values.format(column=column_name_str)
+        else:
+            return column_str_without_values.format(column=column_name_str)
+
+    tables = [
+        table_str.format(
+            table=table_name.lower() if normalize_query else table_name,
+            columns=column_sep.join(
+                map(
+                    lambda y: get_column_str(
+                        table_name=table_name, column_name=y[1]),
+                    filter(
+                        lambda y: y[0] == table_id,
+                        zip(
+                            db_column_names['table_id'],
+                            db_column_names['column_name'],
+                        ),
+                    ),
+                )),
+        ) for table_id, table_name in enumerate(db_table_names)
+    ]
+    if schema_serialization_randomized:
+        random.shuffle(tables)
+    if schema_serialization_with_db_id:
+        serialized_schema = db_id_str.format(
+            db_id=db_id) + table_sep.join(tables)
+    else:
+        serialized_schema = table_sep.join(tables)
+    return serialized_schema
+
+
+def form_input_for_construction(query, question, db_id, db_path, schema):
+    return {
+        'query':
+        query,
+        'question':
+        question,
+        'db_id':
+        db_id,
+        'db_path':
+        db_path,
+        'db_table_names':
+        schema['table_names_original'],
+        'db_column_names': {
+            'table_id': [
+                table_id
+                for table_id, column_name in schema['column_names_original']
+            ],
+            'column_name': [
+                column_name
+                for table_id, column_name in schema['column_names_original']
+            ]
+        },
+        'db_column_types':
+        schema['column_types'],
+        'db_primary_keys': [{
+            'column_id': column_id
+        } for column_id in schema['primary_keys']],
+        'db_foreign_keys': {
+            'column_id': [
+                column_id
+                for column_id, other_column_id in schema['foreign_keys']
+            ],
+            'other_column_id': [
+                other_column_id
+                for column_id, other_column_id in schema['foreign_keys']
+            ]
+        },
+    }

modelscope/preprocessors/ofa/utils/bridge_content_encoder.py

@@ -0,0 +1,266 @@
+"""
+ Copyright (c) 2020, salesforce.com, inc.
+ All rights reserved.
+ SPDX-License-Identifier: BSD-3-Clause
+ For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+
+ Encode DB content.
+"""
+
+import difflib
+import functools
+import sqlite3
+from typing import List, Optional, Tuple
+
+from rapidfuzz import fuzz
+
+# fmt: off
+_stopwords = {
+    'who', 'ourselves', 'down', 'only', 'were', 'him', 'at', "weren't", 'has',
+    'few', "it's", 'm', 'again', 'd', 'haven', 'been', 'other', 'we', 'an',
+    'own', 'doing', 'ma', 'hers', 'all', "haven't", 'in', 'but', "shouldn't",
+    'does', 'out', 'aren', 'you', "you'd", 'himself', "isn't", 'most', 'y',
+    'below', 'is', "wasn't", 'hasn', 'them', 'wouldn', 'against', 'this',
+    'about', 'there', 'don', "that'll", 'a', 'being', 'with', 'your', 'theirs',
+    'its', 'any', 'why', 'now', 'during', 'weren', 'if', 'should', 'those',
+    'be', 'they', 'o', 't', 'of', 'or', 'me', 'i', 'some', 'her', 'do', 'will',
+    'yours', 'for', 'mightn', 'nor', 'needn', 'the', 'until', "couldn't", 'he',
+    'which', 'yourself', 'to', "needn't", "you're", 'because', 'their',
+    'where', 'it', "didn't", 've', 'whom', "should've", 'can', "shan't", 'on',
+    'had', 'have', 'myself', 'am', "don't", 'under', 'was', "won't", 'these',
+    'so', 'as', 'after', 'above', 'each', 'ours', 'hadn', 'having', 'wasn',
+    's', 'doesn', "hadn't", 'than', 'by', 'that', 'both', 'herself', 'his',
+    "wouldn't", 'into', "doesn't", 'before', 'my', 'won', 'more', 'are',
+    'through', 'same', 'how', 'what', 'over', 'll', 'yourselves', 'up',
+    'mustn', "mustn't", "she's", 're', 'such', 'didn', "you'll", 'shan',
+    'when', "you've", 'themselves', "mightn't", 'she', 'from', 'isn', 'ain',
+    'between', 'once', 'here', 'shouldn', 'our', 'and', 'not', 'too', 'very',
+    'further', 'while', 'off', 'couldn', "hasn't", 'itself', 'then', 'did',
+    'just', "aren't"
+}
+# fmt: on
+
+_commonwords = {'no', 'yes', 'many'}
+
+
+def is_number(s: str) -> bool:
+    try:
+        float(s.replace(',', ''))
+        return True
+    except ValueError:
+        return False
+
+
+def is_stopword(s: str) -> bool:
+    return s.strip() in _stopwords
+
+
+def is_commonword(s: str) -> bool:
+    return s.strip() in _commonwords
+
+
+def is_common_db_term(s: str) -> bool:
+    return s.strip() in ['id']
+
+
+class Match(object):
+
+    def __init__(self, start: int, size: int) -> None:
+        self.start = start
+        self.size = size
+
+
+def is_span_separator(c: str) -> bool:
+    return c in "'\"()`,.?! "
+
+
+def split(s: str) -> List[str]:
+    return [c.lower() for c in s.strip()]
+
+
+def prefix_match(s1: str, s2: str) -> bool:
+    i, j = 0, 0
+    for i in range(len(s1)):
+        if not is_span_separator(s1[i]):
+            break
+    for j in range(len(s2)):
+        if not is_span_separator(s2[j]):
+            break
+    if i < len(s1) and j < len(s2):
+        return s1[i] == s2[j]
+    elif i >= len(s1) and j >= len(s2):
+        return True
+    else:
+        return False
+
+
+def get_effective_match_source(s: str, start: int, end: int) -> Match:
+    _start = -1
+
+    for i in range(start, start - 2, -1):
+        if i < 0:
+            _start = i + 1
+            break
+        if is_span_separator(s[i]):
+            _start = i
+            break
+
+    if _start < 0:
+        return None
+
+    _end = -1
+    for i in range(end - 1, end + 3):
+        if i >= len(s):
+            _end = i - 1
+            break
+        if is_span_separator(s[i]):
+            _end = i
+            break
+
+    if _end < 0:
+        return None
+
+    while _start < len(s) and is_span_separator(s[_start]):
+        _start += 1
+    while _end >= 0 and is_span_separator(s[_end]):
+        _end -= 1
+
+    return Match(_start, _end - _start + 1)
+
+
+def get_matched_entries(
+    s: str,
+    field_values: List[str],
+    m_theta: float = 0.85,
+    s_theta: float = 0.85
+) -> Optional[List[Tuple[str, Tuple[str, str, float, float, int]]]]:
+    if not field_values:
+        return None
+
+    if isinstance(s, str):
+        n_grams = split(s)
+    else:
+        n_grams = s
+
+    matched = dict()
+    for field_value in field_values:
+        if not isinstance(field_value, str):
+            continue
+        fv_tokens = split(field_value)
+        sm = difflib.SequenceMatcher(None, n_grams, fv_tokens)
+        match = sm.find_longest_match(0, len(n_grams), 0, len(fv_tokens))
+        if match.size > 0:
+            source_match = get_effective_match_source(n_grams, match.a,
+                                                      match.a + match.size)
+            if source_match and source_match.size > 1:
+                match_str = field_value[match.b:match.b + match.size]
+                source_match_str = s[source_match.start:source_match.start
+                                     + source_match.size]
+                c_match_str = match_str.lower().strip()
+                c_source_match_str = source_match_str.lower().strip()
+                c_field_value = field_value.lower().strip()
+                if (c_match_str and not is_number(c_match_str)
+                        and not is_common_db_term(c_match_str)):
+                    if (is_stopword(c_match_str)
+                            or is_stopword(c_source_match_str)
+                            or is_stopword(c_field_value)):
+                        continue
+                    if c_source_match_str.endswith(c_match_str + "'s"):
+                        match_score = 1.0
+                    else:
+                        if prefix_match(c_field_value, c_source_match_str):
+                            match_score = (
+                                fuzz.ratio(c_field_value, c_source_match_str)
+                                / 100)
+                        else:
+                            match_score = 0
+                    if (is_commonword(c_match_str)
+                            or is_commonword(c_source_match_str)
+                            or is_commonword(c_field_value)
+                        ) and match_score < 1:  # noqa
+                        continue
+                    s_match_score = match_score
+                    if match_score >= m_theta and s_match_score >= s_theta:
+                        if field_value.isupper(
+                        ) and match_score * s_match_score < 1:
+                            continue
+                        matched[match_str] = (
+                            field_value,
+                            source_match_str,
+                            match_score,
+                            s_match_score,
+                            match.size,
+                        )
+
+    if not matched:
+        return None
+    else:
+        return sorted(
+            matched.items(),
+            key=lambda x: (1e16 * x[1][2] + 1e8 * x[1][3] + x[1][4]),
+            reverse=True,
+        )
+
+
+@functools.lru_cache(maxsize=1000, typed=False)
+def get_column_picklist(table_name: str, column_name: str,
+                        db_path: str) -> list:
+    fetch_sql = 'SELECT DISTINCT `{}` FROM `{}`'.format(
+        column_name, table_name)
+    try:
+        conn = sqlite3.connect(db_path)
+        conn.text_factory = bytes
+        c = conn.cursor()
+        c.execute(fetch_sql)
+        picklist = set()
+        for x in c.fetchall():
+            if isinstance(x[0], str):
+                picklist.add(x[0].encode('utf-8'))
+            elif isinstance(x[0], bytes):
+                try:
+                    picklist.add(x[0].decode('utf-8'))
+                except UnicodeDecodeError:
+                    picklist.add(x[0].decode('latin-1'))
+            else:
+                picklist.add(x[0])
+        picklist = list(picklist)
+    finally:
+        conn.close()
+    return picklist
+
+
+def get_database_matches(
+    question: str,
+    table_name: str,
+    column_name: str,
+    db_path: str,
+    top_k_matches: int = 2,
+    match_threshold: float = 0.85,
+) -> List[str]:
+    picklist = get_column_picklist(
+        table_name=table_name, column_name=column_name, db_path=db_path)
+    matches = []
+    if picklist and isinstance(picklist[0], str):
+        matched_entries = get_matched_entries(
+            s=question,
+            field_values=picklist,
+            m_theta=match_threshold,
+            s_theta=match_threshold,
+        )
+        if matched_entries:
+            num_values_inserted = 0
+            for _match_str, (
+                    field_value,
+                    _s_match_str,
+                    match_score,
+                    s_match_score,
+                    _match_size,
+            ) in matched_entries:
+                if 'name' in column_name and match_score * s_match_score < 1:
+                    continue
+                if table_name != 'sqlite_sequence':  # Spider database artifact
+                    matches.append(field_value)
+                    num_values_inserted += 1
+                    if num_values_inserted >= top_k_matches:
+                        break
+    return matches

modelscope/preprocessors/ofa/utils/collate.py

@@ -25,7 +25,7 @@ def collate_fn(samples, pad_idx, eos_idx):
     if samples[0].get('source', None) is not None:
         batch['net_input']['input_ids'] = merge('source')
     if samples[0].get('id', None) is not None:
-        batch['id'] = np.array([s.get['id'] for s in samples])
+        batch['id'] = np.array([s.get('id') for s in samples])
     if samples[0].get('target', None) is not None:
         batch['target'] = merge('target')
         tgt_lengths = torch.LongTensor(
@@ -91,6 +91,20 @@ def collate_fn(samples, pad_idx, eos_idx):
         batch['phone_length'] = torch.tensor(
             [s['phone_target'].size(0) for s in samples], dtype=torch.long)
 
+    # for sudoku
+    if samples[0].get('db_struct', None) is not None:
+        db_struct = [sample['db_struct'] for sample in samples]
+        batch['db_struct'] = db_struct
+    if samples[0].get('mask_ratio', None) is not None:
+        mask_ratio = [sample['mask_ratio'] for sample in samples]
+        batch['mask_ratio'] = mask_ratio
+    if samples[0].get('seg_col_tokens', None) is not None:
+        seg_col_tokens = merge('seg_col_tokens')
+        batch['net_input']['seg_col_tokens'] = seg_col_tokens
+    if samples[0].get('seg_row_tokens', None) is not None:
+        seg_row_tokens = merge('seg_row_tokens')
+        batch['net_input']['seg_row_tokens'] = seg_row_tokens
+
     return batch
 
 

modelscope/preprocessors/ofa/utils/constant.py

@@ -1,3 +1,5 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
 from modelscope.utils.constant import Tasks
 
 OFA_TASK_KEY_MAPPING = {
@@ -11,4 +13,6 @@ OFA_TASK_KEY_MAPPING = {
     Tasks.visual_entailment: ['image', 'text', 'text2'],
     Tasks.text_to_image_synthesis: ['text'],
     Tasks.auto_speech_recognition: ['wav', 'text'],
+    Tasks.sudoku: ['text'],
+    Tasks.text2sql: ['text', 'database'],
 }

modelscope/preprocessors/ofa/utils/get_tables.py

@@ -0,0 +1,88 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+import sqlite3
+import sys
+import traceback
+
+EXIST = {'atis', 'geo', 'advising', 'yelp', 'restaurants', 'imdb', 'academic'}
+
+
+def convert_fk_index(data):
+    fk_holder = []
+    for fk in data['foreign_keys']:
+        tn, col, ref_tn, ref_col = fk[0][0], fk[0][1], fk[1][0], fk[1][1]
+        ref_cid, cid = None, None
+        try:
+            tid = data['table_names_original'].index(tn)
+            ref_tid = data['table_names_original'].index(ref_tn)
+
+            for i, (tab_id,
+                    col_org) in enumerate(data['column_names_original']):
+                if tab_id == ref_tid and ref_col == col_org:
+                    ref_cid = i
+                elif tid == tab_id and col == col_org:
+                    cid = i
+            if ref_cid and cid:
+                fk_holder.append([cid, ref_cid])
+        except ValueError:
+            traceback.print_exc()
+            print('table_names_original: ', data['table_names_original'])
+            print('finding tab name: ', tn, ref_tn)
+            sys.exit()
+    return fk_holder
+
+
+def dump_db_json_schema(db, f):
+    """read table and column info"""
+    conn = sqlite3.connect(db)
+    conn.execute('pragma foreign_keys=ON')
+    cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
+
+    data = {
+        'db_id': f,
+        'table_names_original': [],
+        'table_names': [],
+        'column_names_original': [(-1, '*')],
+        'column_names': [(-1, '*')],
+        'column_types': ['text'],
+        'primary_keys': [],
+        'foreign_keys': [],
+    }
+
+    fk_holder = []
+    for i, item in enumerate(cursor.fetchall()):
+        table_name = item[0]
+        data['table_names_original'].append(table_name)
+        data['table_names'].append(table_name.lower().replace('_', ' '))
+        fks = conn.execute(
+            "PRAGMA foreign_key_list('{}') ".format(table_name)).fetchall()
+        # print("db:{} table:{} fks:{}".format(f,table_name,fks))
+        fk_holder.extend([[(table_name, fk[3]), (fk[2], fk[4])] for fk in fks])
+        cur = conn.execute("PRAGMA table_info('{}') ".format(table_name))
+        for j, col in enumerate(cur.fetchall()):
+            data['column_names_original'].append((i, col[1]))
+            data['column_names'].append((i, col[1].lower().replace('_', ' ')))
+            # varchar, '' -> text, int, numeric -> integer,
+            col_type = col[2].lower()
+            if ('char' in col_type or col_type == '' or 'text' in col_type
+                    or 'var' in col_type):
+                data['column_types'].append('text')
+            elif ('int' in col_type or 'numeric' in col_type
+                  or 'decimal' in col_type or 'number' in col_type
+                  or 'id' in col_type or 'real' in col_type
+                  or 'double' in col_type or 'float' in col_type):
+                data['column_types'].append('number')
+            elif 'date' in col_type or 'time' in col_type or 'year' in col_type:
+                data['column_types'].append('time')
+            elif 'boolean' in col_type:
+                data['column_types'].append('boolean')
+            else:
+                data['column_types'].append('others')
+
+            if col[5] == 1:
+                data['primary_keys'].append(len(data['column_names']) - 1)
+
+    data['foreign_keys'] = fk_holder
+    data['foreign_keys'] = convert_fk_index(data)
+
+    return data

modelscope/utils/constant.py

@@ -152,6 +152,8 @@ class NLPTasks(object):
     extractive_summarization = 'extractive-summarization'
     feature_extraction = 'feature-extraction'
     translation_evaluation = 'translation-evaluation'
+    sudoku = 'sudoku'
+    text2sql = 'text2sql'
 
 
 class AudioTasks(object):

requirements/multi-modal.txt

@@ -1,9 +1,11 @@
 ftfy>=6.0.3
 librosa
+opencv-python
 pycocoevalcap>=1.2
 pycocotools>=2.0.4
 # compatible with taming-transformers-rom1504
 pytorch_lightning<=1.7.7
+rapidfuzz
 # rough-score was just recently updated from 0.0.4 to 0.0.7
 # which introduced compatability issues that are being investigated
 rouge_score<=0.0.4

tests/pipelines/test_ofa_tasks.py

@@ -329,6 +329,43 @@ class OfaTasksTest(unittest.TestCase, DemoCompatibilityCheck):
         for r in result:
             print(r[OutputKeys.TEXT])
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_sudoku_with_name(self):
+        model = 'damo/ofa_sudoku_kaggle_large'
+        ofa_pipe = pipeline(Tasks.sudoku, model=model)
+        # the valid num is 1-9，and use 0 represents the empty block
+        # the separator of column is ` : `, and the separator of row is ` | `
+        example = '5 : 3 : 0 : 0 : 7 : 0 : 0 : 0 : 0 | \
+                6 : 0 : 0 : 1 : 9 : 5 : 0 : 0 : 0 | \
+                0 : 9 : 8 : 0 : 0 : 0 : 0 : 6 : 0 | \
+                8 : 0 : 0 : 0 : 6 : 0 : 0 : 0 : 3 | \
+                4 : 0 : 0 : 8 : 0 : 3 : 0 : 0 : 1 | \
+                7 : 0 : 0 : 0 : 2 : 0 : 0 : 0 : 6 | \
+                0 : 6 : 0 : 0 : 0 : 0 : 2 : 8 : 0 | \
+                0 : 0 : 0 : 4 : 1 : 9 : 0 : 0 : 5 | \
+                0 : 0 : 0 : 0 : 8 : 0 : 0 : 7 : 9'
+
+        result = ofa_pipe(example)
+        print(result[OutputKeys.TEXT])
+        # test batch infer
+        result = ofa_pipe([example for _ in range(3)], batch_size=2)
+        for r in result:
+            print(r[OutputKeys.TEXT])
+
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_text2sql_with_name(self):
+        model = 'damo/ofa_text2sql_spider_large_en'
+        ofa_pipe = pipeline(Tasks.text2sql, model=model)
+        text = 'Show all book categories and the number of books in each category.'
+        database = 'culture_company'  # optional, default `culture_company`
+        example = {'text': text, 'database': database}
+        result = ofa_pipe(example)
+        print(result[OutputKeys.TEXT])
+        # test batch infer
+        result = ofa_pipe([example for _ in range(3)], batch_size=2)
+        for r in result:
+            print(r[OutputKeys.TEXT])
+
     @unittest.skip('demo compatibility test is only enabled on a needed-basis')
     def test_demo_compatibility(self):
         self.compatibility_check()