如何基于HuggingFace构建多Transformer模型的NL2SQL任务流水线？

问题背景 #

需要构建自然语言转SQL（NL2SQL）流水线，涉及实体抽取、查询生成等多个Transformer模型，但HuggingFace默认Pipeline仅支持单模型，想知道规范的构建方法，以及自定义MultiplePipeline是否可行。

一、自定义Pipeline类的可行性 #

你写的自定义Pipeline思路是可行的，但需要补全细节以贴合HuggingFace Pipeline的规范，比如正确继承父类、搭配对应分词器、处理数据流转逻辑。修正后的示例代码如下：

from transformers import Pipeline, AutoModelForTokenClassification, AutoTokenizer, T5ForConditionalGeneration

class NL2SQLPipeline(Pipeline):
    def __init__(self, extractor_model_name, generator_model_name, device=-1, **kwargs):
        # 调用父类初始化方法，指定运行设备
        super().__init__(device=device, **kwargs)
        
        # 加载实体抽取模型及对应分词器
        self.extractor_tokenizer = AutoTokenizer.from_pretrained(extractor_model_name)
        self.extractor = AutoModelForTokenClassification.from_pretrained(extractor_model_name).to(self.device)
        
        # 加载查询生成模型及对应分词器
        self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name)
        self.query_generator = T5ForConditionalGeneration.from_pretrained(generator_model_name).to(self.device)

    def _sanitize_parameters(self, **kwargs):
        # 处理额外参数，无特殊需求返回空字典即可
        return {}, {}, {}

    def _forward(self, inputs):
        # 第一步：执行实体抽取
        extractor_inputs = self.extractor_tokenizer(inputs, return_tensors="pt", padding=True, truncation=True).to(self.device)
        entity_outputs = self.extractor(**extractor_inputs)
        # 解析实体抽取结果（需根据模型输出格式调整）
        entities = self._parse_entities(entity_outputs, extractor_inputs)
        
        # 第二步：构造查询生成的输入文本
        generator_input = f"生成SQL：问句={inputs}，实体={entities}"
        generator_inputs = self.generator_tokenizer(generator_input, return_tensors="pt", padding=True, truncation=True).to(self.device)
        
        # 第三步：生成SQL语句
        sql_outputs = self.query_generator.generate(**generator_inputs)
        sql = self.generator_tokenizer.decode(sql_outputs[0], skip_special_tokens=True)
        
        return {"entities": entities, "sql": sql}

    def _parse_entities(self, outputs, inputs):
        # 从模型输出中提取实体文本和类型
        predictions = outputs.logits.argmax(dim=-1)
        labels = self.extractor.config.id2label
        tokens = self.extractor_tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
        
        entities = []
        current_entity = []
        current_label = None
        for token, pred in zip(tokens, predictions[0]):
            label = labels[pred.item()]
            if label.startswith("B-"):
                if current_entity:
                    entities.append({"text": self.extractor_tokenizer.convert_tokens_to_string(current_entity), "type": current_label})
                    current_entity = []
                current_label = label[2:]
                current_entity.append(token)
            elif label.startswith("I-") and label[2:] == current_label:
                current_entity.append(token)
            else:
                if current_entity:
                    entities.append({"text": self.extractor_tokenizer.convert_tokens_to_string(current_entity), "type": current_label})
                    current_entity = []
                    current_label = None
        if current_entity:
            entities.append({"text": self.extractor_tokenizer.convert_tokens_to_string(current_entity), "type": current_label})
        return entities

二、其他可选方案 #

• 组合独立Pipeline：分别创建实体抽取和查询生成的独立Pipeline，在业务代码中串联调用，这种方式更灵活，无需自定义类，适合快速验证：

  from transformers import pipeline
  
  # 初始化两个独立的Pipeline
  entity_extractor = pipeline("token-classification", model="你的实体抽取模型名称")
  sql_generator = pipeline("text2text-generation", model="你的SQL生成模型名称")
  
  # 定义流水线逻辑
  def nl2sql_pipeline(input_text):
      entities = entity_extractor(input_text)
      # 构造生成模型的输入
      gen_input = f"生成SQL：问句={input_text}，实体={entities}"
      sql = sql_generator(gen_input)[0]["generated_text"]
      return {"entities": entities, "sql": sql}
  

• 端到端封装模型：如果需要多模型联合训练，可以将多个Transformer模型封装成一个nn.Module，用HuggingFace的Trainer进行训练和推理，适合需要联合优化的场景。

内容的提问来源于stack exchange，提问作者zzaebok