app.py

import gradio as gr
from transformers import AutoModelForCausalLM
from transformers import BloomTokenizerFast
from transformers import pipeline, set_seed
import random

model_name = "bloom-560m"
model = AutoModelForCausalLM.from_pretrained(f'jslin09/{model_name}-finetuned-fraud')
tokenizer = BloomTokenizerFast.from_pretrained(f'bigscience/{model_name}', bos_token = '<s>', eos_token = '</s>', pad_token = '<pad>')

def rnd_generate(prompt):
    rnd_seed = random.randint(10, 500)
    set_seed(rnd_seed)
    inputs = tokenizer(prompt, return_tensors="pt")
    results = model.generate(inputs["input_ids"],
                       max_length=500,
                       num_return_sequences=1, 
                       do_sample=True,
                       temperature=0.75,
                       top_k=50, 
                       top_p=0.9)
    return tokenizer.decode(results[0])

def generate(prompt):
    result_length = len(prompt) + 4
    inputs = tokenizer(prompt, return_tensors="pt")
    results = model.generate(inputs["input_ids"],
                       num_return_sequences=2, 
                       max_length=result_length,
                       early_stopping=True,
                       do_sample=True, 
                       top_k=50, 
                       top_p=0.9
                      )
    return tokenizer.decode(results[0])

examples = [
    ["Joaquim sabia que o saldo da conta bancária, o cartão de débito e a senha eram ferramentas financeiras importantes para seu uso pessoal, mas ainda assim as usou de forma fraudulenta."],
    ["Carlos sabia que não tinha condições de pagar pelos gastos no hotel e também não tinha intenção de pagar, mas ainda assim teve a intenção de se apropriar de forma ilícita."],
    ["Paulo teve a intenção de se apropriar de forma ilícita, com base na intenção de cometer furto."],
    ["Durante o tempo em que participou da organização criminosa de fraude, João e membros desconhecidos da organização tiveram a intenção de se apropriar de forma ilícita, cometendo falsificação de documentos privados, falsificação de documentos especiais, fraude e lavagem de dinheiro."],
    ["Com a intenção de cometer falsificação de documentos especiais e fraude, Rafael..."],
    ["Pedro foi anteriormente acusado de fraude..."],
    ["João sabia que recentemente se tornou comum usar números de telefone móveis fictícios, alugados, emprestados ou comprados para fins fraudulentos, sabendo também que o número de telefone móvel é uma ferramenta de comunicação importante na sociedade moderna."]
]

prompts = [
    ["Digite uma frase para a escrita do livro, para que o computador gere a próxima frase. Ou pressione os botões de frases de exemplo abaixo."],
    ["Digite a frase inicial para a escrita do livro, para que o computador gere um rascunho inteiro aleatório."]
]

with gr.Blocks() as demo:
    gr.Markdown(
    """
    <h1 style="text-align: center;">Redação de Documentos Legais</h1>
    """)
    with gr.Row() as row:
        with gr.Column():
            gr.Markdown("""
            <h3>Abstract</h3>
            <p>
            With the development of large-scale language model technology, fine-tuning pre-trained large-scale language models has become a mainstream paradigm to solve downstream tasks of natural language processing. However, training a language model in the legal field requires a large number of legal documents so that the language model can learn legal terminology and the particularity of the format of legal documents. The typical NLP method usually needs to rely on a large number of manually annotation data sets for training. However, in the application of the legal field, it is actually difficult to obtain a large number of manually annotation data sets, which restricted the typical method applied to the task of drafting legal documents. The experimental results of this paper show that not only can a large number of unlabeled legal documents that do not require Chinese word segmentation, but more importantly, it can fine-tune a large pre-trained language model on the local computer to achieve the generating legal document drafts task, and at the same time achieve the protection of information privacy and to improve information security issues.
            </p>
            <h3>Resumo</h3>
            <p>
            Com o desenvolvimento da tecnologia de modelos de linguagem em grande escala, o ajuste fino de modelos de linguagem pré-treinados em grande escala tornou-se um paradigma mainstream para resolver tarefas de processamento de linguagem natural a jusante. No entanto, treinar um modelo de linguagem na área jurídica requer um grande número de documentos legais para que o modelo de linguagem possa aprender a terminologia jurídica e a particularidade do formato dos documentos legais. O método NLP típico geralmente precisa depender de um grande número de conjuntos de dados anotados manualmente para treinamento. No entanto, na aplicação da área jurídica, é realmente difícil obter um grande número de conjuntos de dados anotados manualmente, o que restringiu o método típico aplicado à tarefa de redação de documentos legais. Os resultados experimentais deste artigo mostram que não apenas um grande número de documentos legais não rotulados que não requerem segmentação de palavras em chinês podem ser usados, mas, mais importante, é possível ajustar um grande modelo de linguagem pré-treinado no computador local para alcançar a tarefa de geração de rascunhos de documentos legais e, ao mesmo tempo, alcançar a proteção da privacidade das informações e melhorar as questões de segurança da informação.
            </p>
            """)
        with gr.Column(scale=1, min_width=600):
            with gr.Tab("Assistente Legal de Escrita"):
                result = gr.components.Textbox(lines=7, label="Assistente de Escrita", show_label=True, placeholder=prompts[0])
                prompt = gr.components.Textbox(lines=2, label="Prompt", placeholder=examples[0], visible=False)
                gr.Examples(examples, label='Exemplos', inputs=[prompt])
                prompt.change(generate, inputs=[prompt], outputs=[result])
                btn = gr.Button("Próxima frase")
                btn.click(generate, inputs=[result], outputs=[result])
            with gr.Tab("Geração Aleatória"):
                result2 = gr.components.Textbox(lines=7, label="Geração Aleatória", show_label=True, value = examples[0][0])
                gr.Examples(examples, label='Exemplos', inputs=[result2])
                rnd_btn = gr.Button("Rascunho Aleatório")
                rnd_btn.click(rnd_generate, inputs=[result2], outputs=[result2])
    
if __name__ == "__main__":
    demo.launch()