Legal Court Judgment Prediction (Portuguese)

document_assembler= nlp.DocumentAssembler() \
    .setInputCols(["text"]) \
    .setOutputCols("document")

tokenizer = nlp.Tokenizer() \
    .setInputCols("document") \
    .setOutputCol("token")

seq_classifier = legal.BertForSequenceClassification.pretrained("legclf_judgment_prediction", "pt", "legal/models") \
    .setInputCols(["document", "token"]) \
    .setOutputCol("class")

pipeline = nlp.Pipeline(stages=[
    document_assembler, 
    tokenizer ,
    seq_classifier 
])

# simple examples
example = spark.createDataFrame([["PENAL. PROCESSO PENAL. APELAÇÃO. HOMICÍDIO QUALIFICADO. ARGUIÇÃO DE NULIDADE EM DECORRÊNCIA DA INSCONSTITUCIONALIDADE DO ARTIGO 457 DO CÓDIGO DE PROCESSO PENAL. AFASTADA. PLEITO DE REDIMENSIONAMENTO DA PENA. DOSIMETRIA QUE MERECE RETOQUES. AFASTADA A VALORAÇÃO DESFAVORÁVEL DAS CIRCUNSTÂNCIAS JUDICIAIS DOS ANTECEDENTES E DA PERSONALIDADE DO AGENTE. MANTIDA A CULPABILIDADE, CIRCUNSTÂNCIAS DO DELITO E CONSEQUÊNCIAS DO CRIME. APELO CONHECIDO E PARCIALMENTE PROVIDO. 1 Não há falar em ocorrência de nulidade não caso concreto, não existindo qualquer inconstitucionalidade em virtude do texto legal do ARTIGO 457 do Código de Processo Penal, não tendo ocorrido o adiamento da sessão do júri em virtude da ausência do acusado, conforme alegado. Pelo contrário, este foi devidamente intimado por edital e, mesmo assim, restou ausente. 2 A justificativa apresentada pelo magistrado singular acerca da culpabilidade considerou a alta reprovabilidade da conduta do réu, em virtude da premeditação e frieza na prática delitiva, considerando que o acusado foi até a casa da vítima com o intuito de ceifar a sua vida, experimentando assim a consequência da transgressão, estando acertada a valoração negativa desta circunstância judicial."]]).toDF("text")

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

# result is a DataFrame
result.select("text", "class.result").show(truncate=100)