Legal Question Answering (RoBerta)

documentAssembler = nlp.MultiDocumentAssembler()\
        .setInputCols(["question", "context"])\
        .setOutputCols(["document_question", "document_context"])

spanClassifier = nlp.RoBertaForQuestionAnswering.pretrained("legqa_roberta","en", "legal/models") \
.setInputCols(["document_question", "document_context"]) \
.setOutputCol("answer") \
.setCaseSensitive(True)

pipeline = nlp.Pipeline().setStages([
documentAssembler,
spanClassifier
])

example = spark.createDataFrame([["Who was subjected to torture?", "The applicant submitted that her husband was subjected to treatment amounting to abuse whilst in the custody of police."]]).toDF("question", "context")

result = pipeline.fit(example).transform(example)

result.select('answer.result').show()