如何在 pyspark Z6A8064B5DF4794555500553 上运行自定义 pytorch model（模块）的推理？

Question

如何在 pyspark Z6A8064B5DF4794555500553 上运行自定义 pytorch model（模块）的推理？ 我有一个使用 pytorch model 的 class：

def get_model_obj(model):
    model = model.module if hasattr(model, "module") else model
    return model

class BertEncoder(nn.Module):
    def __init__(
        self, bert_model, output_dim, layer_pulled=-1, add_linear=None):
        super(BertEncoder, self).__init__()
        self.layer_pulled = layer_pulled
        bert_output_dim = bert_model.embeddings.word_embeddings.weight.size(1)

        self.bert_model = bert_model
        if add_linear:
            self.additional_linear = nn.Linear(bert_output_dim, output_dim)
            self.dropout = nn.Dropout(0.1)
        else:
            self.additional_linear = None

    def forward(self, token_ids, segment_ids, attention_mask):
        output_bert, output_pooler = self.bert_model(
            token_ids, segment_ids, attention_mask
        )
        # get embedding of [CLS] token
        if self.additional_linear is not None:
            embeddings = output_pooler
        else:
            embeddings = output_bert[:, 0, :]

        # in case of dimensionality reduction
        if self.additional_linear is not None:
            result = self.additional_linear(self.dropout(embeddings))
        else:
            result = embeddings

        return result

class BiEncoderModule(torch.nn.Module):
    def __init__(self, params):
        super(BiEncoderModule, self).__init__()
        bert_model_path = "bert-base-cased"
        ctxt_bert = BertModel.from_pretrained(bert_model_path)
        cand_bert = BertModel.from_pretrained(bert_model_path)

        self.context_encoder = BertEncoder(
            ctxt_bert,
            params["out_dim"],
            layer_pulled=params["pull_from_layer"],
            add_linear=params["add_linear"],
        )
        self.cand_encoder = BertEncoder(
            cand_bert,
            params["out_dim"],
            layer_pulled=params["pull_from_layer"],
            add_linear=params["add_linear"],
        )
        
        self.config = ctxt_bert.config

    def forward(
            self,
            token_idx_ctxt,
            segment_idx_ctxt,
            mask_ctxt,
            token_idx_cands,
            segment_idx_cands,
            mask_cands,
    ):
        embedding_ctxt = None
        if token_idx_ctxt is not None:
            embedding_ctxt = self.context_encoder(
                token_idx_ctxt, segment_idx_ctxt, mask_ctxt
            )
        embedding_cands = None
        if token_idx_cands is not None:
            embedding_cands = self.cand_encoder(
                token_idx_cands, segment_idx_cands, mask_cands
            )
        return embedding_ctxt, embedding_cands

我可以将它用作简单的 pytorch model：

tokenizer = BertTokenizer.from_pretrained(
            "bert-base-cased" )




def token_to_id(x, max_seq_length):
  cand_tokens = tokenizer.tokenize(x)
  input_ids = tokenizer.convert_tokens_to_ids(cand_tokens)
  padding = [0] * (max_seq_length - len(input_ids))
  input_ids += padding
  return input_ids

 config = {
        "out_dim": 100,
        "pull_from_layer":-1,
        "add_linear" : False
    }
args = argparse.Namespace(**config)
torch_model = BiEncoderModule(params)
ctx_token = token_to_id("[CLS] My name is [unused0] John [unused1] ! [SEP]",32)
ctx_segment =  [0]*len(ctx_token)
ctx_mask = [1 if i !=0 else 0 for i in ctx_token]
cnd_token = token_to_id("[CLS] John [unused3] male given name [SEP]",512)
cnd_segment =  [0]*len(cnd_token)
cndmask = [1 if i !=0 else 0 for i in cnd_token]
embedding_ctxt, embedding_cands = torch_model(torch.as_tensor(ctx_token)[None, :],torch.as_tensor(ctx_segment)[None, :],torch.as_tensor(ctx_mask)[None, :],torch.as_tensor(cnd_token)[None, :],torch.as_tensor(cnd_segment)[None, :],torch.as_tensor(cndmask)[None, :])
embedding_ctxt = embedding_ctxt.unsqueeze(1)  
embedding_cands = embedding_cands.unsqueeze(2)
scores = torch.bmm(embedding_ctxt, embedding_cands)  
scores = torch.squeeze(scores)
print(scores)

但是我怎样才能在 pyspark dataframe 中植入那个 madule？ 我尝试这段代码，但得到“'list' object has no attribute 'toArray'”的错误：

data=[("[CLS] My name is [unused0] John [unused1] ! [SEP]","[CLS] John [unused3] male given name [SEP]")]
spark = sparknlp.start()
df=spark.createDataFrame(data).toDF("context","candidate")

df =df.withColumn("ctx_idx", udf(lambda x : [i for i  in token_to_id(x,32)], ArrayType(IntegerType()))("context"))
df =df.withColumn("ctx_segment", udf(lambda x : [0]*len(x), ArrayType(IntegerType()))("ctx_idx"))
df =df.withColumn("ctx_mask", udf(lambda x :[1 if i !=0 else 0 for i in x], ArrayType(IntegerType()))("ctx_idx"))


df =df.withColumn("cnd_idx", udf(lambda x : [i for i  in token_to_id(x,512)], ArrayType(IntegerType()))("candidate"))
df =df.withColumn("score", udf(lambda ctx_ids, cnd_ids : run_model(ctx_ids, cnd_ids), ArrayType(IntegerType()))("context","candidate"))
df =df.withColumn("cnd_segment", udf(lambda x : [0]*len(x), ArrayType(IntegerType()))("cnd_idx"))
df =df.withColumn("cnd_mask", udf(lambda x : [1 if i !=0 else 0 for i in x], ArrayType(IntegerType()))("cnd_idx"))

df =df.withColumn("ctx_idx", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("ctx_idx"))
df =df.withColumn("ctx_segment", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("ctx_segment"))
df =df.withColumn("ctx_mask", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("ctx_mask"))
df =df.withColumn("cnd_idx", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("cnd_idx"))
df =df.withColumn("cnd_segment", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("cnd_segment"))
df =df.withColumn("cnd_mask", udf(lambda x : [x], ArrayType(ArrayType(IntegerType())))("cnd_mask"))
df = df.withColumn("features", concat(df.ctx_idx,df.ctx_segment, df.ctx_mask, df.cnd_idx, df.cnd_segment,df.cnd_mask))
spark_torch_model = create_spark_torch_model(
    torch_model, 
    inputCol='features',
    predictionCol='predictions'
)

pipeline = Pipeline(stages=[spark_torch_model])
result = pipeline.fit(df).transform(df)
result.printSchema()
result.select("predictions").show()

有什么办法可以完成这项工作吗？ 如何在 pyspark Z6A8064B5DF4794555500553 上运行自定义 pytorch model（模块）的推理？