Update app.py

app.py

@@ -1,22 +1,31 @@
 import time
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
 
 model_id = "nicholasKluge/Aira-Instruct-PT-560M"
+rewardmodel_id = "nicholasKluge/RewardModelPT"
+toxicitymodel_id = "nicholasKluge/ToxicityModelPT"
 token = "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"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+rewardModel = AutoModelForSequenceClassification.from_pretrained(rewardmodel_id, use_auth_token=token)
+toxicityModel = AutoModelForSequenceClassification.from_pretrained(toxicitymodel_id, use_auth_token=token)
 
-if device == "cuda":
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, load_in_8bit=True)
-    
-else:
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+model.eval()
+rewardModel.eval()
+toxicityModel.eval()
 
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
 model.to(device)
+rewardModel.to(device)
+toxicityModel.to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
+rewardTokenizer = AutoTokenizer.from_pretrained(rewardmodel_id, use_auth_token=token)
+toxiciyTokenizer = AutoTokenizer.from_pretrained(toxicitymodel_id, use_auth_token=token)
+
 
 intro = """
 ## O que é `Aira`?
@@ -25,11 +34,19 @@ intro = """
 
 Desenvolvemos os nossos chatbots de conversação de domínio aberto através da geração de texto condicional/ajuste fino por instruções. Esta abordagem tem muitas limitações. Apesar de podermos criar um chatbot capaz de responder a perguntas sobre qualquer assunto, é difícil forçar o modelo a produzir respostas de boa qualidade. E por boa, queremos dizer texto **factual** e **não tóxico**. Isto leva-nos a dois dos problemas mais comuns quando lidando com modelos generativos utilizados em aplicações de conversação:
 
+## Limitações
+
 🤥 Modelos generativos podem perpetuar a geração de conteúdo pseudo-informativo, ou seja, informações falsas que podem parecer verdadeiras.
  
 🤬 Em certos tipos de tarefas, modelos generativos podem produzir conteúdo prejudicial e discriminatório inspirado em estereótipos históricos.
- 
+
+## Uso Intendido
+
 `Aira` destina-se apenas à investigação académica. Para mais informações, visite o nosso [HuggingFace models](https://huggingface.co/nicholasKluge) para ver como desenvolvemos `Aira`.
+
+## Como essa demo funciona?
+
+Esta demonstração utiliza um [`modelo de recompensa`](https://huggingface.co/nicholasKluge/RewardModel) e um [`modelo de toxicidade`](https://huggingface.co/nicholasKluge/ToxicityModel) para avaliar a pontuação de cada resposta candidata, considerando o seu alinhamento com a mensagem do utilizador e o seu nível de toxicidade. A função de geração organiza as respostas candidatas por ordem da sua pontuação de recompensa e elimina as respostas consideradas tóxicas ou nocivas. Posteriormente, a função de geração devolve a resposta candidata com a pontuação mais elevada que ultrapassa o limiar de segurança, ou uma mensagem pré-estabelecida se não forem identificados candidatos seguros.
 """
 
 disclaimer = """
@@ -39,19 +56,20 @@ Se desejar apresentar uma reclamação sobre qualquer mensagem produzida por `Ai
 """
 
 with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
- 
-    gr.Markdown("""<h1><center>Aira Demo (Portuguese) 🤓💬</h1></center>""")
+
+    gr.Markdown("""<h1><center>Aira Demo 🤓💬</h1></center>""")
     gr.Markdown(intro)
-    
+
     chatbot = gr.Chatbot(label="Aira").style(height=500)
+    msg = gr.Textbox(label="Write a question or comment to Aira ...", placeholder="Hi Aira, how are you?")
 
-    with gr.Accordion(label="Parâmetros ⚙️", open=False):
-        top_k = gr.Slider( minimum=10, maximum=100, value=15, step=5, interactive=True, label="Top-k",)
-        top_p = gr.Slider( minimum=0.1, maximum=1.0, value=0.15, step=0.05, interactive=True, label="Top-p",)
-        temperature = gr.Slider( minimum=0.001, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperatura",)
-        max_length = gr.Slider( minimum=10, maximum=500, value=100, step=10, interactive=True, label="Comprimento Máximo",)
-    
-    msg = gr.Textbox(label="Escreva uma pergunta para Aira ...", placeholder="Olá Aira, como você vai?")
+    with gr.Accordion(label="Parâmetros ⚙️", open=True):
+        safety = gr.Radio(["On", "Off"], label="Proteção 🛡️", value="On", info="Ajuda a prevenir o modelo de gerar conteúdo tóxico.")
+        top_k = gr.Slider(minimum=10, maximum=100, value=50, step=5, interactive=True, label="Top-k", info="Controla o número de tokens de maior probabilidade a considerar em cada passo.")
+        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.50, step=0.05, interactive=True, label="Top-p", info="Controla a probabilidade cumulativa dos tokens gerados.")
+        temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperatura", info="Controla a aleatoriedade dos tokens gerados.")
+        max_length = gr.Slider(minimum=10, maximum=500, value=100, step=10, interactive=True, label="Comprimento Máximo", info="Controla o comprimento máximo do texto gerado.")
+        smaple_from = gr.Slider(minimum=2, maximum=10, value=2, step=1, interactive=True, label="Amostragem por Rejeição", info="Controla o número de gerações a partir das quais o modelo de recompensa irá selecionar.")
 
     clear = gr.Button("Limpar Conversa 🧹")
     gr.Markdown(disclaimer)
@@ -59,23 +77,66 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
     def user(user_message, chat_history):
         return gr.update(value=user_message, interactive=True), chat_history + [["👤 " + user_message, None]]
 
-    def generate_response(user_msg, top_p, temperature, top_k, max_length, chat_history):
+    def generate_response(user_msg, top_p, temperature, top_k, max_length, smaple_from, safety, chat_history):
 
-        inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.eos_token, return_tensors="pt").to(device)
+        inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.eos_token, return_tensors="pt").to(model.device)
 
         generated_response = model.generate(**inputs,
             bos_token_id=tokenizer.bos_token_id,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
+            repetition_penalty=1.8,
             do_sample=True,
-            early_stopping=True,   
-            top_k=top_k, 
+            early_stopping=True,
+            top_k=top_k,
             max_length=max_length,
             top_p=top_p,
-            temperature=temperature, 
-            num_return_sequences=1)
+            temperature=temperature,
+            num_return_sequences=smaple_from)
+
+        decoded_text = [tokenizer.decode(tokens, skip_special_tokens=True).replace(user_msg, "") for tokens in generated_response]
+
+        rewards = list()
+        toxicities = list()
+
+        for text in decoded_text:
+          reward_tokens = rewardTokenizer(user_msg, text,
+                        truncation=True,
+                        max_length=512,
+                        return_token_type_ids=False,
+                        return_tensors="pt",
+                        return_attention_mask=True)
+
+          reward_tokens.to(rewardModel.device)
+
+          reward = rewardModel(**reward_tokens)[0].item()
+
+          toxicity_tokens = toxiciyTokenizer(user_msg + " " + text,
+                        truncation=True,
+                        max_length=512,
+                        return_token_type_ids=False,
+                        return_tensors="pt",
+                        return_attention_mask=True)
+
+          toxicity_tokens.to(toxicityModel.device)
+
+          toxicity = toxicityModel(**toxicity_tokens)[0].item()
+
+          rewards.append(reward)
+          toxicities.append(toxicity)
+        
+        toxicity_threshold = 5
+
+        ordered_generations = sorted(zip(decoded_text, rewards, toxicities), key=lambda x: x[1], reverse=True)
+
+        if safety == "On":
+            ordered_generations = [(x, y, z) for (x, y, z) in ordered_generations if z >= toxicity_threshold]
+
+        if len(ordered_generations) == 0:
+          bot_message = """Peço desculpa pelo incómodo, mas parece que não foi possível identificar respostas adequadas que cumpram as nossas normas de segurança. Infelizmente, isto indica que o conteúdo gerado pode conter elementos de toxicidade ou pode não ajudar a responder à sua mensagem. A sua opinião é valiosa para nós e esforçamo-nos por garantir uma conversa segura e construtiva. Não hesite em fornecer mais pormenores ou colocar quaisquer outras questões, e farei o meu melhor para o ajudar."""
 
-        bot_message = tokenizer.decode(generated_response[0], skip_special_tokens=True).replace(user_msg, "")
+        else:
+          bot_message = ordered_generations[0][0]
 
         chat_history[-1][1] = "🤖 "
         for character in bot_message:
@@ -84,10 +145,10 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
             yield chat_history
 
     response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        generate_response, [msg, top_p, temperature, top_k, max_length, chatbot], chatbot
+        generate_response, [msg, top_p, temperature, top_k, max_length, smaple_from, safety, chatbot], chatbot
     )
     response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
-    msg.submit(lambda x: gr.update(value=''), [],[msg])
+    msg.submit(lambda x: gr.update(value=''), None,[msg])
     clear.click(lambda: None, None, chatbot, queue=False)
 
 demo.queue()